Fix static slice indexing with explicit start/stop bounds

stack-info: PR: #440, branch: yf225/stack/56

Author: yf225

helion/_compiler/indexing_strategy.py

@@ -517,7 +517,6 @@ def compute_shape(
                 size = input_size.popleft()
                 # Handle slices with steps
                 slice_size = compute_slice_size(k, size)
-
                 if slice_size != 1:
                     rdim = env.allocate_reduction_dimension(slice_size)
                     output_size.append(rdim.var)
@@ -633,16 +632,24 @@ def create(
                     else:
                         index_values.append(f"{start}{expand}")
                 else:
-                    # Full slice or slice without step
-                    if size != 1:
-                        rdim = env.allocate_reduction_dimension(size)
+                    # Handle slices with start/stop but no step
+                    start = k.start if k.start is not None else 0
+                    stop = k.stop if k.stop is not None else size
+                    slice_size = stop - start
+                    
+                    if slice_size != 1:
+                        rdim = env.allocate_reduction_dimension(slice_size)
                         block_idx = rdim.block_id
                         index_var = state.codegen.index_var(block_idx)
-                        index_values.append(f"({index_var}){expand}")
+                        # Generate index: start + index_var
+                        if start != 0:
+                            index_values.append(f"({start} + ({index_var})){expand}")
+                        else:
+                            index_values.append(f"({index_var}){expand}")
                         if mask := state.codegen.mask_var(block_idx):
                             mask_values.setdefault(f"({mask}){expand}")
                     else:
-                        index_values.append(f"tl.zeros([1], {dtype}){expand}")
+                        index_values.append(f"{start}{expand}")
                 output_idx += 1
             elif isinstance(k, torch.Tensor) and k.ndim == 1:
                 expand = tile_strategy.expand_str(output_size, output_idx)
@@ -941,8 +948,24 @@ def create(
                     res.offsets.append(state.codegen.offset_var(rdim.block_id))
                     res.block_shape.append(rdim.var)
                 else:
-                    res.offsets.append("0")
-                    res.block_shape.append(1)
+                    # Handle slices with start/stop but no step
+                    start = k.start if k.start is not None else 0
+                    stop = k.stop if k.stop is not None else size
+                    slice_size = stop - start
+                    
+                    if slice_size != 1:
+                        env = CompileEnvironment.current()
+                        rdim = env.allocate_reduction_dimension(slice_size)
+                        offset = state.codegen.offset_var(rdim.block_id)
+                        # Add start offset if needed
+                        if start != 0:
+                            res.offsets.append(f"({start} + {offset})")
+                        else:
+                            res.offsets.append(offset)
+                        res.block_shape.append(rdim.var)
+                    else:
+                        res.offsets.append(str(start))
+                        res.block_shape.append(1)
             else:
                 raise exc.InvalidIndexingType(k)
         res.validate()

helion/_compiler/type_propagation.py

@@ -467,7 +467,6 @@ def _device_indexing_size(self, key: TypeInfo) -> list[int | torch.SymInt]:
 
                 # For slices with steps, we need to calculate the output size differently
                 output_size = compute_slice_size(slice_obj, size)
-
                 if self.origin.is_device():
                     output_sizes.append(output_size)
                 elif output_size != 1:
@@ -510,8 +509,9 @@ def propagate_setitem(
             lhs_rank = len(lhs_shape)
             if isinstance(value, TensorType):
                 rhs_rank = value.fake_value.ndim
-                # Allow scalar tensors (rank 0) to be assigned to any rank (broadcasts)
-                if rhs_rank != 0 and lhs_rank != rhs_rank:
+                rhs_numel = value.fake_value.numel()
+                # Allow scalar tensors (rank 0) or single-element tensors to be assigned to any rank (broadcasts)
+                if rhs_rank != 0 and rhs_numel != 1 and lhs_rank != rhs_rank:
                     raise exc.RankMismatch(
                         lhs_rank,
                         rhs_rank,

helion/_compiler/utils.py

@@ -25,5 +25,8 @@ def compute_slice_size(
         stop = slice_obj.stop if slice_obj.stop is not None else original_size
         step = slice_obj.step
         return (stop - start + step - 1) // step
-    # Full slice or slice without step
-    return original_size
+    else:
+        # Calculate slice size based on start/stop
+        start = slice_obj.start if slice_obj.start is not None else 0
+        stop = slice_obj.stop if slice_obj.stop is not None else original_size
+        return stop - start

test/test_indexing.py

@@ -877,7 +877,6 @@ def kernel(
         torch.testing.assert_close(src_result, expected_src)
         torch.testing.assert_close(dst_result, expected_dst)
 
-    @skipIfNormalMode("InternalError: Unexpected type <class 'slice'>")
     def test_range_slice(self):
         """Test both setter from scalar and getter for [10:20]"""
 
@@ -903,9 +902,7 @@ def kernel(
         torch.testing.assert_close(src_result, expected_src)
         torch.testing.assert_close(dst_result, expected_dst)
 
-    @skipIfNormalMode(
-        "InternalError: AssertionError in type_propagation.py - slice indexing error"
-    )
+    @skipIfNormalMode("Dynamic slices (i:i+1) are not supported - FX cannot trace symbolic slice indices")
     def test_range_slice_dynamic(self):
         """Test both [i:i+1] = scalar and [i] = [i:i+1] patterns"""