tensor[tile] when tile size is 1 returns a 1D tensor, instead of a scalar

Author: joydddd

helion/_compiler/indexing_strategy.py

@@ -378,9 +378,8 @@ def create(
         assert len(index_values) == fake_value.ndim
         index_expr = []
         for i, idx in enumerate(index_values):
-            if fake_value.size(i) != 1:
-                stride = state.device_function.tensor_stride(fake_value, i).name
-                index_expr.append(f"{idx} * {stride}")
+            stride = state.device_function.tensor_stride(fake_value, i).name
+            index_expr.append(f"{idx} * {stride}")
         if not index_expr:
             shape_str = tile_strategy.shape_str(output_size)
             index_expr.append(f"tl.zeros({shape_str}, {dtype})")

helion/_compiler/tile_strategy.py

@@ -670,11 +670,19 @@ def codegen_device_loop(self, state: CodegenState) -> DeviceLoopState:
                 type_comment=None,
             )
             assert for_node.body is body
-            extra_body = [
-                statement_from_string(
-                    f"{index_var} = {offset_var} + tl.arange(0, ({block_size_var})).to({dtype})"
-                ),
-            ]
+            extra_body = []
+            if block_size == 1:
+                extra_body.append(
+                    statement_from_string(
+                        f"{index_var} = {offset_var} + tl.zeros([1], {dtype})"
+                    ),
+                )
+            else:
+                extra_body.append(
+                    statement_from_string(
+                        f"{index_var} = {offset_var} + tl.arange(0, ({block_size_var})).to({dtype})"
+                    ),
+                )
             mask_statement = self._setup_mask(  # pyright: ignore[reportAttributeAccessIssue]
                 state, block_idx, block_size, index_var, end
             )