jax-hpc-profiler 0.2.3__py3-none-any.whl → 0.2.6__py3-none-any.whl

jax_hpc_profiler/create_argparse.py

@@ -136,23 +136,25 @@ def create_argparser():
         default="%m%-%f%-%pn%-%pr%-%b%-%p%-%n%")
 
     args = parser.parse_args()
-
-    if 'plot_all' in args.pdim_strategy and len(args.pdim_strategy) > 1:
-        print(
-            "Warning: 'plot_all' strategy is combined with other strategies. Using 'plot_all' only."
-        )
-        args.pdim_strategy = ['plot_all']
-
-    if 'plot_fastest' in args.pdim_strategy and len(args.pdim_strategy) > 1:
-        print(
-            "Warning: 'plot_fastest' strategy is combined with other strategies. Using 'plot_fastest' only."
-        )
-        args.pdim_strategy = ['plot_fastest']
-    if args.plot_times is not None:
-        args.plot_columns = args.plot_times
-    elif args.plot_memory is not None:
-        args.plot_columns = args.plot_memory
-    else:
-        raise ValueError('Either plot_times or plot_memory should be provided')
+    
+    # if command was plot, then check if pdim_strategy is validat
+    if args.command == 'plot':
+      if 'plot_all' in args.pdim_strategy and len(args.pdim_strategy) > 1:
+          print(
+              "Warning: 'plot_all' strategy is combined with other strategies. Using 'plot_all' only."
+          )
+          args.pdim_strategy = ['plot_all']
+
+      if 'plot_fastest' in args.pdim_strategy and len(args.pdim_strategy) > 1:
+          print(
+              "Warning: 'plot_fastest' strategy is combined with other strategies. Using 'plot_fastest' only."
+          )
+          args.pdim_strategy = ['plot_fastest']
+      if args.plot_times is not None:
+          args.plot_columns = args.plot_times
+      elif args.plot_memory is not None:
+          args.plot_columns = args.plot_memory
+      else:
+          raise ValueError('Either plot_times or plot_memory should be provided')
 
     return args

jax_hpc_profiler/timer.py

@@ -28,15 +28,14 @@ class Timer:
             return None
         return cost_analysis[0]['flops']
 
-    def _normalize_memory_units(self, memory_analysis) -> str:  
-        
-        sizes_str = ['B', 'KB', 'MB', 'GB', 'TB' , 'PB']
-        factors = [1 , 1024 , 1024**2 , 1024**3 , 1024**4 , 1024**5]
+    def _normalize_memory_units(self, memory_analysis) -> str:
+
+        sizes_str = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
+        factors = [1, 1024, 1024**2, 1024**3, 1024**4, 1024**5]
         factor = int(np.log10(memory_analysis) // 3)
-        
+
         return f"{memory_analysis / factors[factor]:.2f} {sizes_str[factor]}"
 
-    
     def _read_memory_analysis(self, memory_analysis: Any) -> Tuple:
         if memory_analysis is None:
             return None, None, None, None
@@ -68,12 +67,9 @@ class Timer:
         self.compiled_code["LOWERED"] = lowered.as_text()
         self.compiled_code["COMPILED"] = compiled.as_text()
         self.profiling_data["FLOPS"] = cost_analysis
-        self.profiling_data[
-            "generated_code"] = memory_analysis[0]
-        self.profiling_data[
-            "argument_size"] = memory_analysis[1]
-        self.profiling_data[
-            "output_size"] = memory_analysis[2]
+        self.profiling_data["generated_code"] = memory_analysis[0]
+        self.profiling_data["argument_size"] = memory_analysis[1]
+        self.profiling_data["output_size"] = memory_analysis[2]
         self.profiling_data["temp_size"] = memory_analysis[3]
         return out
 
@@ -101,7 +97,7 @@ class Timer:
         global_times = jax.make_array_from_callback(
             shape=global_shape,
             sharding=sharding,
-            data_callback=lambda _: jnp.expand_dims(times_array,axis=0))
+            data_callback=lambda _: jnp.expand_dims(times_array, axis=0))
 
         @partial(shard_map,
                  mesh=mesh,
@@ -141,90 +137,90 @@ class Timer:
         z = x if z is None else z
 
         times_array = self._get_mean_times()
-        min_time = np.min(times_array)
-        max_time = np.max(times_array)
-        mean_time = np.mean(times_array)
-        std_time = np.std(times_array)
-        last_time = times_array[-1]
-
-
-        flops = self.profiling_data["FLOPS"]
-        generated_code = self.profiling_data["generated_code"]
-        argument_size = self.profiling_data["argument_size"]
-        output_size = self.profiling_data["output_size"]
-        temp_size = self.profiling_data["temp_size"]
-
-        csv_line = (
-            f"{function},{precision},{x},{y},{z},{px},{py},{backend},{nodes},"
-            f"{self.jit_time:.4f},{min_time:.4f},{max_time:.4f},{mean_time:.4f},{std_time:.4f},{last_time:.4f},"
-            f"{generated_code},{argument_size},{output_size},{temp_size},{flops}\n"
-        )
-
-        with open(csv_filename, 'a') as f:
-            f.write(csv_line)
-
-        param_dict = {
-            "Function": function,
-            "Precision": precision,
-            "X": x,
-            "Y": y,
-            "Z": z,
-            "PX": px,
-            "PY": py,
-            "Backend": backend,
-            "Nodes": nodes,
-        }
-        param_dict.update(extra_info)
-        profiling_result = {
-            "JIT Time": self.jit_time,
-            "Min Time": min_time,
-            "Max Time": max_time,
-            "Mean Time": mean_time,
-            "Std Time": std_time,
-            "Last Time": last_time,
-            "Generated Code": generated_code,
-            "Argument Size": argument_size,
-            "Output Size": output_size,
-            "Temporary Size": temp_size,
-            "FLOPS": self.profiling_data["FLOPS"]
-        }
-        iteration_runs = {}
-        for i in range(len(times_array)):
-            iteration_runs[f"Run {i}"] = times_array[i]
-
-        with open(md_filename, 'w') as f:
-            f.write(f"# Reporting for {function}\n")
-            f.write(f"## Parameters\n")
-            f.write(
-                tabulate(param_dict.items(),
-                         headers=["Parameter", "Value"],
-                         tablefmt='github'))
-            f.write("\n---\n")
-            f.write(f"## Profiling Data\n")
-            f.write(
-                tabulate(profiling_result.items(),
-                         headers=["Parameter", "Value"],
-                         tablefmt='github'))
-            f.write("\n---\n")
-            f.write(f"## Iteration Runs\n")
-            f.write(
-                tabulate(iteration_runs.items(),
-                         headers=["Iteration", "Time"],
-                         tablefmt='github'))
-            f.write("\n---\n")
-            f.write(f"## Compiled Code\n")
-            f.write(f"```hlo\n")
-            f.write(self.compiled_code["COMPILED"])
-            f.write(f"\n```\n")
-            f.write("\n---\n")
-            f.write(f"## Lowered Code\n")
-            f.write(f"```hlo\n")
-            f.write(self.compiled_code["LOWERED"])
-            f.write(f"\n```\n")
-            f.write("\n---\n")
-            if self.save_jaxpr:
-                f.write(f"## JAXPR\n")
-                f.write(f"```haskel\n")
-                f.write(self.compiled_code["JAXPR"])
+        if jax.process_index() == 0:
+
+            min_time = np.min(times_array)
+            max_time = np.max(times_array)
+            mean_time = np.mean(times_array)
+            std_time = np.std(times_array)
+            last_time = times_array[-1]
+
+            flops = self.profiling_data["FLOPS"]
+            generated_code = self.profiling_data["generated_code"]
+            argument_size = self.profiling_data["argument_size"]
+            output_size = self.profiling_data["output_size"]
+            temp_size = self.profiling_data["temp_size"]
+
+            csv_line = (
+                f"{function},{precision},{x},{y},{z},{px},{py},{backend},{nodes},"
+                f"{self.jit_time:.4f},{min_time:.4f},{max_time:.4f},{mean_time:.4f},{std_time:.4f},{last_time:.4f},"
+                f"{generated_code},{argument_size},{output_size},{temp_size},{flops}\n"
+            )
+
+            with open(csv_filename, 'a') as f:
+                f.write(csv_line)
+
+            param_dict = {
+                "Function": function,
+                "Precision": precision,
+                "X": x,
+                "Y": y,
+                "Z": z,
+                "PX": px,
+                "PY": py,
+                "Backend": backend,
+                "Nodes": nodes,
+            }
+            param_dict.update(extra_info)
+            profiling_result = {
+                "JIT Time": self.jit_time,
+                "Min Time": min_time,
+                "Max Time": max_time,
+                "Mean Time": mean_time,
+                "Std Time": std_time,
+                "Last Time": last_time,
+                "Generated Code": self._normalize_memory_units(generated_code),
+                "Argument Size": self._normalize_memory_units(argument_size),
+                "Output Size": self._normalize_memory_units(output_size),
+                "Temporary Size": self._normalize_memory_units(temp_size),
+                "FLOPS": self.profiling_data["FLOPS"]
+            }
+            iteration_runs = {}
+            for i in range(len(times_array)):
+                iteration_runs[f"Run {i}"] = times_array[i]
+
+            with open(md_filename, 'w') as f:
+                f.write(f"# Reporting for {function}\n")
+                f.write(f"## Parameters\n")
+                f.write(
+                    tabulate(param_dict.items(),
+                             headers=["Parameter", "Value"],
+                             tablefmt='github'))
+                f.write("\n---\n")
+                f.write(f"## Profiling Data\n")
+                f.write(
+                    tabulate(profiling_result.items(),
+                             headers=["Parameter", "Value"],
+                             tablefmt='github'))
+                f.write("\n---\n")
+                f.write(f"## Iteration Runs\n")
+                f.write(
+                    tabulate(iteration_runs.items(),
+                             headers=["Iteration", "Time"],
+                             tablefmt='github'))
+                f.write("\n---\n")
+                f.write(f"## Compiled Code\n")
+                f.write(f"```hlo\n")
+                f.write(self.compiled_code["COMPILED"])
                 f.write(f"\n```\n")
-
+                f.write("\n---\n")
+                f.write(f"## Lowered Code\n")
+                f.write(f"```hlo\n")
+                f.write(self.compiled_code["LOWERED"])
+                f.write(f"\n```\n")
+                f.write("\n---\n")
+                if self.save_jaxpr:
+                    f.write(f"## JAXPR\n")
+                    f.write(f"```haskel\n")
+                    f.write(self.compiled_code["JAXPR"])
+                    f.write(f"\n```\n")

{jax_hpc_profiler-0.2.3.dist-info → jax_hpc_profiler-0.2.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jax_hpc_profiler
-Version: 0.2.3
+Version: 0.2.6
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License: GNU GENERAL PUBLIC LICENSE

jax_hpc_profiler-0.2.6.dist-info/RECORD

@@ -0,0 +1,12 @@
+jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
+jax_hpc_profiler/create_argparse.py,sha256=dEicamRYqJ6GGdgcph2bwAbmdxPkS4tS12xZ4c0X_Pk,6484
+jax_hpc_profiler/main.py,sha256=CKsKVUKsMRatlYfrFLOV1WZ582rZPtofV89sY_2tpQI,2370
+jax_hpc_profiler/plotting.py,sha256=iA-CCBsDK5wM872tmpS6qo4ws9F6MZLutfkdAN_sWlw,8326
+jax_hpc_profiler/timer.py,sha256=j6oH5IZz12VJik2cE7EQ3a9tAW9C8xl7D2QLW8Bkz3s,8617
+jax_hpc_profiler/utils.py,sha256=GqlGaD-Zf9GmtuQRJAF6S6QD4E2S8QtZ-jGjjC5ZFwU,14133
+jax_hpc_profiler-0.2.6.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jax_hpc_profiler-0.2.6.dist-info/METADATA,sha256=AgLXyb89gdxgeyDv02_P5oqIvoffdTh1mc3zFPUVBAU,49250
+jax_hpc_profiler-0.2.6.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+jax_hpc_profiler-0.2.6.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
+jax_hpc_profiler-0.2.6.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
+jax_hpc_profiler-0.2.6.dist-info/RECORD,,

jax_hpc_profiler-0.2.3.dist-info/RECORD

@@ -1,12 +0,0 @@
-jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
-jax_hpc_profiler/create_argparse.py,sha256=sY3OKe6lMrXtVnKyx-EtREXLy9L1TK_mdf0WYRQXu5A,6351
-jax_hpc_profiler/main.py,sha256=CKsKVUKsMRatlYfrFLOV1WZ582rZPtofV89sY_2tpQI,2370
-jax_hpc_profiler/plotting.py,sha256=iA-CCBsDK5wM872tmpS6qo4ws9F6MZLutfkdAN_sWlw,8326
-jax_hpc_profiler/timer.py,sha256=baE5DRsQBYRBiphkceTi4qI_8FPGKQEh73f2pAeS-oc,8208
-jax_hpc_profiler/utils.py,sha256=GqlGaD-Zf9GmtuQRJAF6S6QD4E2S8QtZ-jGjjC5ZFwU,14133
-jax_hpc_profiler-0.2.3.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.3.dist-info/METADATA,sha256=myC-zD7y_pRb_-tZoSFi0KmglZH8Gk88_-U5RE14Q04,49250
-jax_hpc_profiler-0.2.3.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-jax_hpc_profiler-0.2.3.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.3.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.3.dist-info/RECORD,,