[Refactor] Integrate env & configs for an unified configuration file

Author: YWHyuk

.github/workflows/pytorchsim_test.yml

@@ -152,7 +152,7 @@ def load(cls, source_code,
         else:
             link_option = ""
         # Generate LLVM kernel calller and binary for validation
-        if extension_config.CONFIG_TORCHSIM_FUNCTIONAL_MODE:
+        if extension_config.pytorchsim_functional_mode:
             # Use custom malloc to avoid size error
             new_link_option = link_option + " -Wl,--wrap=malloc -Wl,--wrap=free"
             cmds = mlir_compile_command(new_input_path, vectorlane_size, vlen=vlen)
@@ -169,7 +169,7 @@ def load(cls, source_code,
                     print("Error output:", e.output)
                     assert(0)
 
-                val_llvm_caller = MLIRKernelCallerCodeGen(extension_config.CONFIG_TORCHSIM_FUNCTIONAL_MODE, arg_attributes)
+                val_llvm_caller = MLIRKernelCallerCodeGen(extension_config.pytorchsim_functional_mode, arg_attributes)
                 val_llvm_caller.generate_wrapper_file(write_path, validation_wrapper_name)
                 val_llvm_caller.compile_wih_kernel(write_path, key, validation_wrapper_name,
                                                    validation_binary_name, new_link_option)
@@ -200,7 +200,7 @@ def load(cls, source_code,
                 print("Error output:", e.output)
                 assert(0)
 
-            if not extension_config.CONFIG_TORCHSIM_TIMING_MODE:
+            if not extension_config.pytorchsim_timing_mode:
                 return key
 
             # Generate MLIR kernel calller and binary for cycle calculation
@@ -271,13 +271,13 @@ def dummy_simulator(*args, **kwargs):
                 # Dump arguments and meta data
                 dump_metadata(args, arg_attributes, result_path)
                 runtime_path = FunctionalSimulator.get_runtime_dump_path(result_path)
-                if not autotune and (extension_config.CONFIG_TORCHSIM_FUNCTIONAL_MODE or validate):
+                if not autotune and (extension_config.pytorchsim_functional_mode or validate):
                     funcsim = FunctionalSimulator(result_path, key)
                     funcsim.run_spike(args, arg_attributes,
                                     runtime_path, self.validation_binary_name,
                                     vectorlane_size=vectorlane_size, spad_info=spad_info,
-                                    cleanup=extension_config.CONFIG_CLEANUP_DUMP_ARGS, silent_mode=silent_mode)
-                if not extension_config.CONFIG_TORCHSIM_TIMING_MODE:
+                                    silent_mode=silent_mode)
+                if not extension_config.pytorchsim_timing_mode:
                     return
 
                 onnx_path = os.path.join(result_path, "tile_graph.onnx")
@@ -303,12 +303,11 @@ def dryrun_simulator(*args, **kwargs):
                 runtime_path = FunctionalSimulator.get_runtime_dump_path(result_path)
 
                 # Todo. Support valude dependent mode for graph mode
-                if False: # extension_config.CONFIG_TORCHSIM_FUNCTIONAL_MODE:
+                if False: # extension_config.pytorchsim_functional_mode:
                     funcsim = FunctionalSimulator(result_path, key)
                     funcsim.run_spike(args, arg_attributes,
                                     runtime_path, self.validation_binary_name,
-                                    vectorlane_size=vectorlane_size, spad_info=spad_info,
-                                    cleanup=extension_config.CONFIG_CLEANUP_DUMP_ARGS)
+                                    vectorlane_size=vectorlane_size, spad_info=spad_info)
             return result_path, runtime_path, None
 
         is_dryrun = int(os.environ.get('TOGSIM_EAGER_MODE', default=False)) and not autotune

PyTorchSimFrontend/extension_codecache.py

@@ -3,112 +3,89 @@
 import tempfile
 import importlib
 import datetime
+import json
+
+CONFIG_TORCHSIM_DIR = os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim')
+CONFIG_GEM5_PATH = os.environ.get('GEM5_PATH', default="/workspace/gem5/build/RISCV/gem5.opt")
+CONFIG_TORCHSIM_LLVM_PATH = os.environ.get('TORCHSIM_LLVM_PATH', default="/usr/bin")
+
+CONFIG_TORCHSIM_DUMP_MLIR_IR = int(os.environ.get("TORCHSIM_DUMP_MLIR_IR", default=False))
+CONFIG_TORCHSIM_DUMP_LLVM_IR = int(os.environ.get("TORCHSIM_DUMP_LLVM_IR", default=False))
 
 def __getattr__(name):
+    # TOGSim config
+    config_path = os.environ.get('TOGSIM_CONFIG',
+                default=f"{CONFIG_TORCHSIM_DIR}/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.json")
+    if name == "CONFIG_TOGSIM_CONFIG":
+        return config_path
+    config_json = json.load(open(config_path, 'r'))
 
     # Hardware info config
-    if name == "CONFIG_VECTOR_LANE":
-        return int(os.environ.get("TORCHSIM_VECTOR_LANE", default=128))
-    if name == "CONFIG_VECTOR_LANE_STRIDE":
-        return int(os.environ.get("TORCHSIM_VECTOR_LANE_STRIDE", default=2))
+    if name == "vpu_num_lanes":
+        return config_json["vpu_num_lanes"]
     if name == "CONFIG_SPAD_INFO":
         return {
           "spad_vaddr" : 0xD0000000,
           "spad_paddr" : 0x2000000000,
-          "spad_size" : int(os.environ.get("TORCHSIM_SPAD_SIZE", default=128)) << 10 # Note: spad size per lane
+          "spad_size" : config_json["vpu_spad_size_kb_per_lane"] << 10 # Note: spad size per lane
         }
+
     if name == "CONFIG_PRECISION":
         return 4 # 32bit
     if name == "CONFIG_NUM_CORES":
-        return 1
-    if name == "CONFIG_VLEN":
-        return 256 # 256bits / 32bits = 8 [elements]
+        return config_json["num_cores"]
+    if name == "vpu_vector_length_bits":
+        return config_json["vpu_vector_length_bits"]
 
-    # Tile size config
-    if name == "CONFIG_TORCHSIM_DIR":
-          return os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim')
+    if name == "pytorchsim_functional_mode":
+        return config_json['pytorchsim_functional_mode']
+    if name == "pytorchsim_timing_mode":
+        return config_json['pytorchsim_timing_mode']
 
-    if name == "CONFIG_TORCHSIM_DUMP_PATH":
-        return os.environ.get('TORCHSIM_DUMP_PATH', default = __getattr__('CONFIG_TORCHSIM_DIR'))
-    if name == "CONFIG_TORCHSIM_LOG_PATH":
-        return os.environ.get('TORCHSIM_DUMP_LOG_PATH', default = os.path.join(__getattr__("CONFIG_TORCHSIM_DIR"), "outputs", datetime.datetime.now().strftime('%Y%m%d_%H%M%S')))
-    if name == "CONFIG_TORCHSIM_FUNCTIONAL_MODE":
-        return int(os.environ.get('TORCHSIM_FUNCTIONAL_MODE', default=True))
-    if name == "CONFIG_TORCHSIM_TIMING_MODE":
-        return int(os.environ.get("TORCHSIM_TIMING_MODE", True))
-    if name == "CONFIG_CLEANUP_DUMP_ARGS":
-        return int(os.environ.get('CLEANUP_DUMP_ARGS', default=False))
-
-    # LLVM PATH
-    if name == "CONFIG_TORCHSIM_LLVM_PATH":
-        return os.environ.get('TORCHSIM_LLVM_PATH', default="/usr/bin")
-    if name == "CONFIG_TORCHSIM_DUMP_MLIR_IR":
-        return int(os.environ.get("TORCHSIM_DUMP_MLIR_IR", default=False))
-    if name == "CONFIG_TORCHSIM_DUMP_LLVM_IR":
-        return int(os.environ.get("TORCHSIM_DUMP_LLVM_IR", default=False))
+    # Mapping strategy
+    if name == "codegen_mapping_strategy":
+        codegen_mapping_strategy = config_json["codegen_mapping_strategy"]
+        assert(codegen_mapping_strategy in ["heuristic", "autotune", "external-then-heuristic", "external-then-autotune"]), "Invalid mapping strategy!"
+        return codegen_mapping_strategy
 
-    # TOGSim config
-    if name == "CONFIG_TOGSIM_CONFIG":
-        return os.environ.get('TOGSIM_CONFIG',
-                default=f"{__getattr__('CONFIG_TORCHSIM_DIR')}/configs/systolic_ws_128x128_c1_simple_noc_tpuv3.json")
-    if name == "CONFIG_TOGSIM_EAGER_MODE":
-        return int(os.environ.get("TOGSIM_EAGER_MODE", default=False))
-    if name == "CONFIG_TOGSIM_DEBUG_LEVEL":
-        return os.environ.get("TOGSIM_DEBUG_LEVEL", "")
-
-    # GEM5 config
-    if name == "CONFIG_GEM5_PATH":
-        return os.environ.get('GEM5_PATH', default="/workspace/gem5/build/RISCV/gem5.opt")
-
-    # Mapping Policy
-    if name == "CONFIG_MAPPING_POLICY":
-        return os.environ.get('TORCHSIM_MAPPING_POLICY', default="heuristic") # heuristic, manual, autotune
-
-    # Manual Tile Size
-    if name == "CONFIG_TILE_M":
-        return int(os.getenv("TORCHSIM_TILE_M", __getattr__("CONFIG_VECTOR_LANE")))
-    if name == "CONFIG_TILE_N":
-        return int(os.getenv("TORCHSIM_TILE_N", __getattr__("CONFIG_VECTOR_LANE")))
-    if name == "CONFIG_TILE_K":
-        return int(os.getenv("TORCHSIM_TILE_K", __getattr__("CONFIG_VECTOR_LANE")))
-
-    if name == "CONFIG_MANUAL_SUBTILE_SIZE":
-        return int(os.environ.get('TORCHSIM_MANUAL_SUBTILE_SIZE', default=False))
-    if name == "CONFIG_SUBTILE_M":
-        return int(os.environ.get('TORCHSIM_SUBTILE_M', default=__getattr__("CONFIG_VECTOR_LANE")))
-    if name == "CONFIG_SUBTILE_N":
-        return int(os.environ.get('TORCHSIM_SUBTILE_N', default=__getattr__("CONFIG_VECTOR_LANE")))
-    if name == "CONFIG_SUBTILE_K":
-        return int(os.environ.get('TORCHSIM_SUBTILE_K', default=__getattr__("CONFIG_VECTOR_LANE")))
+    if name == "codegen_external_mapping_file":
+        return config_json["codegen_external_mapping_file"]
 
     # Autotune config
-    if name == "CONFIG_MAX_AUTOTUNE_TRY":
-        return int(os.environ.get('MAX_AUTOTUNE_TRY', default=10))
-    if name == "CONFIG_AUTOTUNE_TEMPLATE_TOPK":
-        return int(os.environ.get('AUTOTUNE_TEMPLATE_TOPK', default=4))
-
-    if name == "CONFIG_GEMM_CHEATSHEET_PATH":
-          return os.environ.get('TORCHSIM_GEMM_CHEATSHEET_PATH',
-                          default=f"{__getattr__('CONFIG_TORCHSIM_DIR')}/validation/gemm_tpuv3_cheatsheet.json")
+    if name == "codegen_autotune_max_retry":
+        return config_json["codegen_autotune_max_retry"]
+    if name == "codegen_autotune_template_topk":
+        return config_json["codegen_autotune_template_topk"]
+
     # Compiler Optimization
-    if name == "CONFIG_COMPILER_OPTIMIZATION":
-        return os.environ.get('TORCHSIM_COMPILER_OPTIMIZATION', default="all")  # options: all, none, custom
+    if name == "codegen_compiler_optimization":
+        return config_json["codegen_compiler_optimization"]
 
     # Advanced fusion options
     if name == "CONFIG_FUSION":
-        return True if (__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all" or "fusion" in __getattr__("CONFIG_COMPILER_OPTIMIZATION")) else False
+        return True if (__getattr__("codegen_compiler_optimization") == "all" or "fusion" in __getattr__("codegen_compiler_optimization")) else False
     if name == "CONFIG_FUSION_REDUCTION_EPILOGUE":
-        return True if (__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all" or "reduction_epliogue" in __getattr__("CONFIG_COMPILER_OPTIMIZATION")) else False
+        return True if (__getattr__("codegen_compiler_optimization") == "all" or "reduction_epliogue" in __getattr__("codegen_compiler_optimization")) else False
     if name == "CONFIG_FUSION_REDUCTION_REDUCTION":
-        return True if (__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all" or "reduction_reduction" in __getattr__("CONFIG_COMPILER_OPTIMIZATION")) else False
+        return True if (__getattr__("codegen_compiler_optimization") == "all" or "reduction_reduction" in __getattr__("codegen_compiler_optimization")) else False
     if name == "CONFIG_FUSION_PROLOGUE":
-        return True if ((__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all") or ("prologue" in __getattr__("CONFIG_COMPILER_OPTIMIZATION"))) else False
+        return True if ((__getattr__("codegen_compiler_optimization") == "all") or ("prologue" in __getattr__("codegen_compiler_optimization"))) else False
     if name == "CONFIG_SINGLE_BATCH_CONV":
-        return True if (__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all" or "single_batch_conv" in __getattr__("CONFIG_COMPILER_OPTIMIZATION")) else False
+        return True if (__getattr__("codegen_compiler_optimization") == "all" or "single_batch_conv" in __getattr__("codegen_compiler_optimization")) else False
     if name == "CONFIG_MULTI_TILE_CONV":
-        return True if (__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all" or "multi_tile_conv" in __getattr__("CONFIG_COMPILER_OPTIMIZATION")) else False
+        return True if (__getattr__("codegen_compiler_optimization") == "all" or "multi_tile_conv" in __getattr__("codegen_compiler_optimization")) else False
     if name == "CONFIG_SUBTILE":
-        return True if (__getattr__("CONFIG_COMPILER_OPTIMIZATION") == "all" or "subtile" in __getattr__("CONFIG_COMPILER_OPTIMIZATION")) else False
+        return True if (__getattr__("codegen_compiler_optimization") == "all" or "subtile" in __getattr__("codegen_compiler_optimization")) else False
+
+    if name == "CONFIG_TOGSIM_DEBUG_LEVEL":
+        return os.environ.get("TOGSIM_DEBUG_LEVEL", "")
+    if name == "CONFIG_TORCHSIM_DUMP_PATH":
+        return os.environ.get('TORCHSIM_DUMP_PATH', default = CONFIG_TORCHSIM_DIR)
+    if name == "CONFIG_TORCHSIM_LOG_PATH":
+        return os.environ.get('TORCHSIM_DUMP_LOG_PATH', default = os.path.join(CONFIG_TORCHSIM_DIR, "outputs", datetime.datetime.now().strftime('%Y%m%d_%H%M%S')))
+
+    if name == "CONFIG_TOGSIM_EAGER_MODE":
+        return int(os.environ.get("TOGSIM_EAGER_MODE", default=False))
 
 # SRAM Buffer allocation plan
 def load_plan_from_module(module_path):

PyTorchSimFrontend/extension_config.py

@@ -1627,7 +1627,7 @@ def make_choices(self, nodes, kernel_name):
     def autotune(self, *args):
         def get_cycle(choice):
             bench_runner = choice[0]
-            for n_try in range(extension_config.CONFIG_MAX_AUTOTUNE_TRY): # TODO: make simple
+            for n_try in range(extension_config.codegen_autotune_max_retry): # TODO: make simple
                 try:
                     out = bench_runner()
                     return out[-1]
@@ -1664,7 +1664,7 @@ def run_bench(self, nodes, kernel_name, src_code):
                 "spad_info": self.spad_info,
                 "vlen" : self.vlen,
                 "arg_attributes" : arg_attributes,
-                "validate" : extension_config.CONFIG_TORCHSIM_FUNCTIONAL_MODE,
+                "validate" : extension_config.pytorchsim_functional_mode,
                 "autotune" : True,
             },
             source_code=src_code,
@@ -1683,7 +1683,7 @@ def _log_autotune_result(self, best_choice, best_cycle):
     def codegen_nodes(self, nodes, kernel_name):
         src_code = super().codegen_nodes(nodes, kernel_name)
         self._prepare_simulator_headers(src_code)
-        if extension_config.CONFIG_MAPPING_POLICY == "autotune" and extension_config.CONFIG_TORCHSIM_TIMING_MODE:
+        if extension_config.CONFIG_MAPPING_POLICY == "autotune" and extension_config.pytorchsim_timing_mode:
             optimal_src_code = self.autotune(nodes, kernel_name)[0]
             if optimal_src_code is not None:
                 return optimal_src_code

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -567,11 +567,11 @@ def set_tile_info(self, tile_desc : MLIRMultiDimTile):
 class BaseMLIRHardwareInfo():
     def __init__(self):
         # Default HW setting
-        self.vector_lane = extension_config.CONFIG_VECTOR_LANE
+        self.vector_lane = extension_config.vpu_num_lanes
         self.spad_info = extension_config.CONFIG_SPAD_INFO
         self.precision = extension_config.CONFIG_PRECISION
         self.num_cores = extension_config.CONFIG_NUM_CORES
-        self.vlen = extension_config.CONFIG_VLEN
+        self.vlen = extension_config.vpu_vector_length_bits
 
 class BaseMLIRKernel(common.Kernel, BaseMLIRHardwareInfo):
     newvar_prefix = "%"
@@ -700,7 +700,7 @@ def extract_dividers(self, implicit_ops):
 
     def compute_tile_size(self, nodes, vars, reduction_vars):
         vlane_split_axis = len(vars) - 1
-        vlane_stride = extension_config.CONFIG_VECTOR_LANE_STRIDE
+        vlane_stride = 2 # Set minimum vlane stride
 
         # Set initial tile size & vector lane mapping
         if self.kernel_group.tile_desc is None:

PyTorchSimFrontend/mlir/mlir_common.py

@@ -93,7 +93,7 @@ def outer_func_render(self, kernel_name, input_args):
             OUTPUT=Y,
             PADDING_H=self.padding[0],
             PADDING_W=self.padding[1],
-            VALIDATION_MODE=extension_config.CONFIG_TORCHSIM_FUNCTIONAL_MODE,
+            VALIDATION_MODE=extension_config.pytorchsim_functional_mode,
             TOGSIM_EAGER_MODE=eager_mode,
             input_reorder=self.input_reorder
         )

PyTorchSimFrontend/mlir/mlir_conv_common.py

@@ -110,7 +110,7 @@ def convolution(
         mlir_template = MLIRConvSingleBatchTemplate([x, weight, bias], layout, **kwargs)
     elif BATCH == 1 and stride[0] != 1 and extension_config.CONFIG_SINGLE_BATCH_CONV:
         mlir_template = MLIRConvSingleBatchStridedTemplate([x, weight, bias], layout, **kwargs)
-    elif I_C < extension_config.CONFIG_VECTOR_LANE // 8 and extension_config.CONFIG_MULTI_TILE_CONV: # 8 is hard-coded for now. This should be changed to a better heuristic.
+    elif I_C < extension_config.vpu_num_lanes // 8 and extension_config.CONFIG_MULTI_TILE_CONV: # 8 is hard-coded for now. This should be changed to a better heuristic.
         mlir_template = MLIRConvMultiTileTemplate([x, weight, bias], layout, **kwargs)
     else:
         mlir_template = MLIRConvTemplate([x, weight, bias], layout, **kwargs)

PyTorchSimFrontend/mlir/mlir_lowering.py

@@ -257,7 +257,7 @@ def define_kernel(self, src_code, kernel_name, vector_lane, spad_info, loop_size
             codecache_def.writeline(f"spad_info={spad_info},")
             codecache_def.writeline(f"origins={origins},")
             codecache_def.writeline("arg_attributes=arg_attributes,")
-            codecache_def.writeline(f"vlen={extension_config.CONFIG_VLEN})")
+            codecache_def.writeline(f"vlen={extension_config.vpu_vector_length_bits})")
             wrapper.define_kernel(kernel_name, codecache_def.getvalue(), cuda=False)
         return kernel_name
 

PyTorchSimFrontend/mlir/mlir_scheduling.py

@@ -1232,7 +1232,7 @@ def make_kernel_render(
                 template=self,
                 kwargs=kwargs
             )
-            tile_candidates = self.get_tile_candidates(**kwargs)[:extension_config.CONFIG_AUTOTUNE_TEMPLATE_TOPK]
+            tile_candidates = self.get_tile_candidates(**kwargs)[:extension_config.codegen_autotune_template_topk]
             return kernel, tile_candidates, render
 
         return MLIRTemplateCaller(

PyTorchSimFrontend/mlir/mlir_template.py

@@ -147,7 +147,7 @@ Simulation consists of three steps
 
 If you want to turn off the `SpikeSimulator` for fast simulation, you can set as below.
 ```bash
-export TORCHSIM_FUNCTIONAL_MODE=False
+export pytorchsim_functional_mode=False
 ```
 Log contains memory & core stats.
 ```bash
@@ -329,8 +329,8 @@ Last but not least, you must set `l2d_type` and `l2d_config` in the [TOGSim conf
 
 You can configure these options using environment variables.
 ```bash
-export TORCHSIM_VECTOR_LANE=128 # vector lane size
-export TORCHSIM_VECTOR_LANE_STRIDE=2  # vector lane stride for DMA
+export vpu_num_lanes=128 # vector lane size
+export vpu_num_lanes_STRIDE=2  # vector lane stride for DMA
 export TORCHSIM_DIR=/workspace/PyTorchSim # home directory
 
 # Plan which tensor allocated in TPUv4's CMEM
@@ -342,7 +342,7 @@ export TORCHSIM_USE_TIMING_POOLING=0 # use lightweight pooling for timing
 ## TOGSim Configuration
 ![NPU_Core](./docs/npu_core.jpg)
 
-`TOGSim/configs` directory contains example NPU configuration files in the JSON format.
+`configs` directory contains example NPU configuration files in the JSON format.
 ```
   "num_cores" : 2,                   // Number of NPU cores
   "core_freq_mhz" : 940,             // Core's frequency (MHz)