[Fix] Keep loop_size variable in the choice list

Author: YWHyuk

.github/workflows/pytorchsim_test.yml

@@ -674,6 +674,9 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
+      - name: Prepare volume directory
+        run: mkdir -p /tmp/torchsim-ci/${GITHUB_SHA}
+
       - name: Run run_cycle.sh
         run: |
           echo "Running run_cycle.sh"
@@ -682,7 +685,9 @@ jobs:
             -e TORCHSIM_DUMP_PATH=/dump \
             -e TORCHSIM_VECTOR_LANE="${{ inputs.vector_lane }}" \
             -e TORCHSIM_SPAD_SIZE="${{ inputs.spad_size }}" \
-            ${{ inputs.image_name }} PyTorchSim/experiments/artifact/cycle_validation/run_cycle.sh && cp /workspace/PyTorchSim/experiments/artifact/cycle_validation/summary_cycle.out /dump/summary_cycle.out
+            ${{ inputs.image_name }} bash -c \
+            "cd /workspace && PyTorchSim/experiments/artifact/cycle_validation/run_cycle.sh && \
+            cp PyTorchSim/experiments/artifact/cycle_validation/summary_cycle.out /dump/summary_cycle.out"
 
       - name: Upload Accuracy Report Artifact
         uses: actions/upload-artifact@v4

PyTorchSimFrontend/mlir/mlir_codegen_backend.py

@@ -1627,15 +1627,15 @@ def get_cycle(choice):
         choices = self.make_choices(*args)
 
         if len(choices) == 0: # can't autotune
-            return None
+            return [None, None]
         with ThreadPoolExecutor(max_workers=8) as executor:
             results = list(executor.map(get_cycle, choices))
         max_idx = results.index(min(results))
         if min(results) == float("inf"):
             raise RuntimeError("Failed to find optimal tile size...")
         self._log_autotune_result(choices[max_idx], results[max_idx])
-        optimal_src_code = choices[max_idx][1]
-        return optimal_src_code
+        optimal_src_code, loop_size = choices[max_idx][1], choices[max_idx][-1]
+        return optimal_src_code, loop_size
 
     def _log_autotune_result(self, best_choice, best_cycle):
         print(
@@ -1648,7 +1648,7 @@ 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_AUTOTUNE and not extension_config.CONFIG_BACKENDSIM_SPIKE_ONLY:
-            optimal_src_code = self.autotune(nodes, kernel_name)
+            optimal_src_code = self.autotune(nodes, kernel_name)[0]
             if optimal_src_code is not None:
                 return optimal_src_code
         return src_code

PyTorchSimFrontend/mlir/mlir_template.py

@@ -29,7 +29,7 @@
 from PyTorchSimFrontend.mlir.mlir_scheduling import SchedulerNode
 from torch._inductor.codegen import common
 
-from PyTorchSimFrontend.extension_config import CONFIG_TORCHSIM_DIR, CONFIG_AUTOTUNE_TEMPLATE_TOPK
+from PyTorchSimFrontend.extension_config import CONFIG_TORCHSIM_DIR, CONFIG_AUTOTUNE_TEMPLATE_TOPK, CONFIG_AUTOTUNE_TEMPLATE
 from . import mlir_common
 
 class IndentedBufferGroup:
@@ -494,7 +494,7 @@ def make_choices(self, tile_candidates, render, template_node, prologue_nodes, e
             print(f"[Auto-tune] Trying tile size: {list(tile_info)}")
             src_code = self.codegen_template_code(render, template_node, prologue_nodes, epilogue_nodes, tile_info)
             bench_runner = self.run_bench([template_node], self.kernel_name, src_code)
-            choices.append((bench_runner, src_code, tile_info))
+            choices.append((bench_runner, src_code, tile_info, self.loop_size))
             self.reset(reason=None)
         return choices
 
@@ -506,7 +506,12 @@ def _log_autotune_result(self, best_choice, best_cycle):
         )
 
     def codegen_nodes(self, tile_candidates, render, template_node, prologue_nodes, epilogue_nodes):
-        src_code = self.autotune(tile_candidates, render, template_node, prologue_nodes, epilogue_nodes)
+        if CONFIG_AUTOTUNE_TEMPLATE and len(tile_candidates):
+            src_code, loop_size = self.autotune(tile_candidates, render, template_node, prologue_nodes, epilogue_nodes)
+            self.loop_size = loop_size
+        else:
+            tile_info = tile_candidates[0] if tile_candidates else None
+            src_code = self.codegen_template_code(render, template_node, prologue_nodes, epilogue_nodes, tile_info)
 
         with V.set_kernel_handler(self):
             self.meta_kernel()