jax-hpc-profiler 0.2.6__py3-none-any.whl → 0.2.8__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,6 +134,8 @@ 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%")
 
+    subparsers.add_parser('label_help',help='Label customization help')
+
     args = parser.parse_args()
     
     # if command was plot, then check if pdim_strategy is validat

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

@@ -7,10 +7,10 @@ import pandas as pd
 from matplotlib.axes import Axes
 from matplotlib.patches import FancyBboxPatch
 
-from .utils import inspect_df, plot_with_pdims_strategy
+from .utils import inspect_df, plot_with_pdims_strategy, inspect_data
 
 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})
 
 
@@ -117,13 +118,13 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
         axs = axs.flatten()
     else:
         axs = [axs]
-
+    
     for i, fixed_size in enumerate(fixed_sizes):
         ax: Axes = axs[i]
 
+        x_values = []
+        y_values = []
         for method, df in dataframes.items():
-            x_values = []
-            y_values = []
 
             filtered_method_df = df[df[fixed_column] == int(fixed_size)]
             if filtered_method_df.empty:
@@ -136,6 +137,7 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
                                    plot_columns)
 
             for backend, precision, function, plot_column in combinations:
+              
                 filtered_params_df = filtered_method_df[
                     (filtered_method_df['backend'] == backend)
                     & (filtered_method_df['precision'] == precision) &
@@ -148,10 +150,11 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
 
                 x_values.extend(x_vals)
                 y_values.extend(y_vals)
-
-        plotting_memory = 'time' not in plot_columns[0].lower()
-        configure_axes(ax, x_values, y_values, f"{title} {fixed_size}", xlabel,
-                       plotting_memory, memory_units)
+        
+        if len(x_values) != 0: 
+          plotting_memory = 'time' not in plot_columns[0].lower()
+          configure_axes(ax, x_values, y_values, f"{title} {fixed_size}", xlabel,
+                         plotting_memory, memory_units)
 
     for i in range(num_subplots, num_rows * num_cols):
         fig.delaxes(axs[i])

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]
@@ -145,7 +138,6 @@ class Timer:
             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"]
@@ -154,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({
@@ -336,12 +336,11 @@ def clean_up_csv(
         # Filter data sizes
         if data_sizes:
             df = df[df['x'].isin(data_sizes)]
-
+        
         # Filter pdims
         if pdims:
             px_list, py_list = zip(*[map(int, p.split('x')) for p in pdims])
             df = df[(df['px'].isin(px_list)) & (df['py'].isin(py_list))]
-
         # convert memory units columns to remquested memory_units
         match memory_units:
             case 'KB':
@@ -385,6 +384,7 @@ def clean_up_csv(
             df.drop(columns=['px', 'py'], inplace=True)
             if not 'plot_all' in pdims_strategy:
                 df = df[df['decomp'].isin(pdims_strategy)]
+
         # check available gpus in dataset
         available_gpu_counts.update(df['gpus'].unique())
         available_data_sizes.update(df['x'].unique())

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

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

jax_hpc_profiler-0.2.8.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=jnWpyfZOB5w5L4BNIU5euKbUt__wfGPyCjUITrEwScM,8431
+jax_hpc_profiler/timer.py,sha256=qPp3NcCJlMM-Cmw2mEWn63BlvPqmj_k7E8P9m0-Fy7k,8294
+jax_hpc_profiler/utils.py,sha256=tUXnNHwQSSCqA6XxLd4MoV2gyFIC7ncB-uvVc2INhms,14140
+jax_hpc_profiler-0.2.8.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jax_hpc_profiler-0.2.8.dist-info/METADATA,sha256=NDOKx7RzKr4Z4mJnY457-4uxaP2TDhwWYH5rZRxfukQ,49250
+jax_hpc_profiler-0.2.8.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
+jax_hpc_profiler-0.2.8.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
+jax_hpc_profiler-0.2.8.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
+jax_hpc_profiler-0.2.8.dist-info/RECORD,,

{jax_hpc_profiler-0.2.6.dist-info → jax_hpc_profiler-0.2.8.dist-info}/WHEEL

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

jax_hpc_profiler-0.2.6.dist-info/RECORD

@@ -1,12 +0,0 @@
-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,,