Add extensive setter/getter unit tests for indexed tensor; fix bugs discovered by new tests (#422)

Author: yf225

helion/_compiler/indexing_strategy.py

@@ -518,7 +518,10 @@ def boundary_check(self, state: CodegenState) -> str:
             return repr(result)
         return "None"
 
-    def need_reshape(self) -> bool:
+    def need_reshape(self, node: ast.AST) -> bool:
+        if isinstance(node, ast.Constant):
+            # Don't reshape scalar constants - they will be broadcast automatically
+            return False
         if len(self.reshaped_size) != len(self.block_shape):
             return True
         env = CompileEnvironment.current()
@@ -528,13 +531,13 @@ def need_reshape(self) -> bool:
         return False
 
     def reshape_load(self, state: CodegenState, node: ast.AST) -> ast.AST:
-        if not self.need_reshape():
+        if not self.need_reshape(node):
             return node
         shape = state.tile_strategy.shape_str(self.reshaped_size)
         return expr_from_string(f"tl.reshape(node, {shape})", node=node)
 
     def reshape_store(self, state: CodegenState, node: ast.AST) -> ast.AST:
-        if not self.need_reshape():
+        if not self.need_reshape(node):
             return node
         shape = state.tile_strategy.shape_str(self.block_shape)
         return expr_from_string(f"tl.reshape(node, {shape})", node=node)

helion/_testing.py

@@ -40,6 +40,11 @@ def skipIfRefEager(reason: str) -> Callable[[Callable], Callable]:
     return unittest.skipIf(os.environ.get("HELION_INTERPRET") == "1", reason)
 
 
+def skipIfNormalMode(reason: str) -> Callable[[Callable], Callable]:
+    """Skip test if running in normal mode (i.e. if HELION_INTERPRET=1 is not set)."""
+    return unittest.skipIf(os.environ.get("HELION_INTERPRET") != "1", reason)
+
+
 @contextlib.contextmanager
 def track_run_ref_calls() -> Generator[list[int], None, None]:
     """Context manager that tracks BoundKernel.run_ref calls.

helion/autotuner/config_generation.py

@@ -52,8 +52,10 @@ def _collect_spec(spec: ConfigSpecFragment) -> object:
             for i, spec in enumerate(self.flat_spec)
             if spec.category() == Category.NUM_WARPS
         )
-        self.min_block_size: int = max(
-            [spec.min_size for spec in config_spec.block_sizes]
+        self.min_block_size: int = (
+            max([spec.min_size for spec in config_spec.block_sizes])
+            if config_spec.block_sizes
+            else 1
         )
 
     def unflatten(self, flat_values: FlatConfig) -> Config:
@@ -80,7 +82,9 @@ def get_next_value(spec: ConfigSpecFragment) -> object:
 
     def block_numel(self, flat_config: FlatConfig) -> int:
         return functools.reduce(
-            operator.mul, [cast("int", flat_config[i]) for i in self.block_size_indices]
+            operator.mul,
+            [cast("int", flat_config[i]) for i in self.block_size_indices],
+            1,
         )
 
     def shrink_config(