Fix negative indexing and multi-dimensional slicing in Helion

stack-info: PR: #438, branch: yf225/stack/54

Author: yf225

helion/_compiler/indexing_strategy.py

@@ -549,7 +549,21 @@ def create(
                         index_values.append(f"tl.zeros([1], {dtype}){expand}")
                     output_idx += 1
             elif isinstance(k, int):
-                index_values.append(repr(k))
+                # Normalize negative indices
+                if k < 0:
+                    dim_idx = len(index_values)
+                    dim_size = fake_value.size(dim_idx)
+                    # Handle both concrete and symbolic dimension sizes
+                    if isinstance(dim_size, int):
+                        normalized_k = k + dim_size
+                        index_values.append(repr(normalized_k))
+                    else:
+                        # For symbolic dimensions, we need to generate the proper expression
+                        # The state.codegen is a GenerateAST instance which has device_function
+                        sympy_expr = dim_size._sympy_() + k
+                        index_values.append(f"({state.codegen.device_function.user_sympy_expr(sympy_expr)})")
+                else:
+                    index_values.append(repr(k))
             elif isinstance(k, torch.SymInt):
                 symbol = k._sympy_()
                 origin = None
@@ -839,7 +853,21 @@ def create(
                         res.offsets.append("0")
                         res.block_shape.append(1)
             elif isinstance(k, int):
-                res.offsets.append(repr(k))
+                # Normalize negative indices
+                if k < 0:
+                    dim_idx = len(res.offsets)
+                    dim_size = fake_value.size(dim_idx)
+                    # Handle both concrete and symbolic dimension sizes
+                    if isinstance(dim_size, int):
+                        normalized_k = k + dim_size
+                        res.offsets.append(repr(normalized_k))
+                    else:
+                        # For symbolic dimensions, we need to generate the proper expression
+                        # The state.codegen is a GenerateAST instance which has device_function
+                        sympy_expr = dim_size._sympy_() + k
+                        res.offsets.append(f"({state.codegen.device_function.user_sympy_expr(sympy_expr)})")
+                else:
+                    res.offsets.append(repr(k))
                 res.block_shape.append(1)
             elif isinstance(k, torch.SymInt):
                 symbol = k._sympy_()

test/test_indexing.py

@@ -733,7 +733,6 @@ def kernel(
         torch.testing.assert_close(src2_result, expected_src2)
         torch.testing.assert_close(dst2_result, expected_dst2)
 
-    @skipIfNormalMode("InternalError: Negative indexes")
     def test_negative_indexing(self):
         """Test both setter from scalar and getter for [-1]"""
 
@@ -784,9 +783,6 @@ def kernel(
         torch.testing.assert_close(src_result, expected_src)
         torch.testing.assert_close(dst_result, expected_dst)
 
-    @skipIfNormalMode(
-        "RankMismatch: Cannot assign a tensor of rank 2 to a buffer of rank 3"
-    )
     def test_multi_dim_slice(self):
         """Test both setter from scalar and getter for [:, :, i]"""