jax-hpc-profiler 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

jax_hpc_profiler/create_argparse.py

@@ -67,7 +67,6 @@ def create_argparser():
     plot_parser.add_argument('-fn',
                              '--function_name',
                              nargs='+',
-                             default=['FFT'],
                              help='Function names to filter')
 
     # Time or memory related arguments
@@ -135,24 +134,28 @@ def create_argparser():
         'Custom label for the plot. You can use placeholders: %decomposition% (or %p%), %precision% (or %pr%), %plot_name% (or %pn%), %backend% (or %b%), %node% (or %n%), %methodname% (or %m%)',
         default="%m%-%f%-%pn%-%pr%-%b%-%p%-%n%")
 
-    args = parser.parse_args()
+    subparsers.add_parser('label_help',help='Label customization help')
 
-    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')
+    args = parser.parse_args()
+    
+    # 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/main.py

@@ -11,6 +11,15 @@ def main():
     if args.command == 'concat':
         input_dir, output_dir = args.input, args.output
         concatenate_csvs(input_dir, output_dir)
+    elif args.command == 'label_help':
+        print(f"Customize the label text for the plot. using these commands.")
+        print(' -- %m% or %methodname%: method name')
+        print(' -- %f% or %function%: function name')
+        print(' -- %pn% or %plot_name%: plot name')
+        print(' -- %pr% or %precision%: precision')
+        print(' -- %b% or %backend%: backend')
+        print(' -- %p% or %pdims%: pdims')
+        print(' -- %n% or %node%: node')
     elif args.command == 'plot':
         dataframes, available_gpu_counts, available_data_sizes = clean_up_csv(
             args.csv_files, args.precision, args.function_name, args.gpus,
@@ -23,11 +32,10 @@ def main():
             f"requested GPUS: {args.gpus} available GPUS: {available_gpu_counts}"
         )
         # filter back the requested data sizes and gpus
-        args.gpus = [gpu for gpu in args.gpus if gpu in available_gpu_counts]
-        args.data_size = [
-            data_size for data_size in args.data_size
-            if data_size in available_data_sizes
-        ]
+        
+        args.gpus = available_gpu_counts if args.gpus is None else [gpu for gpu in args.gpus if gpu in available_gpu_counts]
+        args.data_size = available_data_sizes if args.data_size is None else [data_size for data_size in args.data_size if data_size in available_data_sizes]
+
         if len(args.gpus) == 0:
             print(f"No dataframes found for the given GPUs. Exiting...")
             sys.exit(1)

jax_hpc_profiler/plotting.py

@@ -10,7 +10,7 @@ from matplotlib.patches import FancyBboxPatch
 from .utils import inspect_df, plot_with_pdims_strategy
 
 np.seterr(divide='ignore')
-plt.rcParams.update({'font.size': 10})
+plt.rcParams.update({'font.size': 15})
 
 
 def configure_axes(ax: Axes,
@@ -38,7 +38,7 @@ def configure_axes(ax: Axes,
     f2 = lambda x: np.log2(x)
     g2 = lambda x: 2**x
     ax.set_xlim([min(x_values), max(x_values)])
-    y_min, y_max = min(y_values) * 0.9, max(y_values) * 1.1
+    y_min, y_max = min(y_values) * 0.6, max(y_values) * 1.1
     ax.set_title(title)
     ax.set_ylim([y_min, y_max])
     ax.set_xscale('function', functions=(f2, g2))
@@ -57,6 +57,7 @@ def configure_axes(ax: Axes,
     ax.legend(loc='lower center',
               bbox_to_anchor=(0.5, 0.05),
               ncol=4,
+              fontsize="x-large",
               prop={'size': 14})
 
 

jax_hpc_profiler/timer.py

@@ -23,16 +23,11 @@ class Timer:
         self.compiled_code = {}
         self.save_jaxpr = save_jaxpr
 
-    def _read_cost_analysis(self, cost_analysis: Any) -> str | None:
-        if cost_analysis is None:
-            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]
-        factor = int(np.log10(memory_analysis) // 3)
+        factor = 0 if memory_analysis == 0 else int(np.log10(memory_analysis) // 3)
 
         return f"{memory_analysis / factors[factor]:.2f} {sizes_str[factor]}"
 
@@ -62,11 +57,9 @@ class Timer:
         compiled = lowered.compile()
         memory_analysis = self._read_memory_analysis(
             compiled.memory_analysis())
-        cost_analysis = self._read_cost_analysis(compiled.cost_analysis())
 
         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]
@@ -138,13 +131,13 @@ class Timer:
 
         times_array = self._get_mean_times()
         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"]
@@ -153,7 +146,7 @@ class Timer:
             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"
+                f"{generated_code},{argument_size},{output_size},{temp_size}\n"
             )
 
             with open(csv_filename, 'a') as f:

jax_hpc_profiler/utils.py

@@ -123,7 +123,7 @@ def plot_with_pdims_strategy(ax: Axes, df: pd.DataFrame, method: str,
         # Sort all and keep fastest
         sorted_dfs = []
         for _, group in df_decomp:
-            group.sort_values(by=[y_col], inplace=True, ascending=False)
+            group.sort_values(by=[y_col], inplace=True, ascending=True)
             sorted_dfs.append(group.iloc[0])
         sorted_df = pd.DataFrame(sorted_dfs)
         label_params.update({

{jax_hpc_profiler-0.2.5.dist-info → jax_hpc_profiler-0.2.7.dist-info}/METADATA

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

jax_hpc_profiler-0.2.7.dist-info/RECORD

@@ -0,0 +1,12 @@
+jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
+jax_hpc_profiler/create_argparse.py,sha256=m9_lg9HHxq2JDMITiHXQW1Ximua0ClwsEq1Zd9Y0hvo,6511
+jax_hpc_profiler/main.py,sha256=VJKvVc4m2XGJI2yp9ZF9tmmBmnTDpZ7-6LGo8ZIrWLc,2906
+jax_hpc_profiler/plotting.py,sha256=cwHznCZ2pF2J7AtyUOB3pASnahKBLRWHAPGXmGDvWas,8360
+jax_hpc_profiler/timer.py,sha256=qPp3NcCJlMM-Cmw2mEWn63BlvPqmj_k7E8P9m0-Fy7k,8294
+jax_hpc_profiler/utils.py,sha256=okWQUJHblUKkYnw7j7wJ75PSbhVItXKkTMKjj0BmgR0,14132
+jax_hpc_profiler-0.2.7.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jax_hpc_profiler-0.2.7.dist-info/METADATA,sha256=bQkpy5Kr8ybEM7GU7qR0FEnDV7xsLbrq98GRDfgDTQU,49250
+jax_hpc_profiler-0.2.7.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
+jax_hpc_profiler-0.2.7.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
+jax_hpc_profiler-0.2.7.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
+jax_hpc_profiler-0.2.7.dist-info/RECORD,,

{jax_hpc_profiler-0.2.5.dist-info → jax_hpc_profiler-0.2.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.1.0)
+Generator: setuptools (73.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

jax_hpc_profiler-0.2.5.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=r4Mw2tC82cxvkMPkIy8BuZjKikgxn6cviEgmu6rpC9o,8616
-jax_hpc_profiler/utils.py,sha256=GqlGaD-Zf9GmtuQRJAF6S6QD4E2S8QtZ-jGjjC5ZFwU,14133
-jax_hpc_profiler-0.2.5.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.5.dist-info/METADATA,sha256=6Vk6fA1nz-m8ZZVzSZPsg9xR9iJkFJ01x36pddG0RAM,49250
-jax_hpc_profiler-0.2.5.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-jax_hpc_profiler-0.2.5.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.5.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.5.dist-info/RECORD,,