Fix tensor value assignment and partial indexing in Helion

Author: yf225

helion/_compiler/indexing_strategy.py

@@ -102,6 +102,25 @@ def codegen_store(
     ) -> ast.AST:
         indexing = SubscriptIndexing.create(state, fake_tensor, subscript, extra_mask)
         name = state.device_function.tensor_arg(fake_tensor).name
+        
+        # Check if value is a tensor load (Name node with id matching a tensor arg)
+        if isinstance(value, ast.Name) and hasattr(state.device_function, '_tensor_args'):
+            # Check if this name corresponds to a tensor argument
+            for tensor, tensor_arg in state.device_function._tensor_args.items():
+                if tensor_arg.name == value.id:
+                    # This is a tensor value, we need to load from it
+                    # Get the shape of the slice we're storing to
+                    output_shape = SubscriptIndexing.compute_shape(fake_tensor, subscript)
+                    if len(output_shape) == 1 and tensor.ndim == 1:
+                        # Load the entire 1D tensor
+                        value_indexing = SubscriptIndexing.create(state, tensor, [slice(None)], None)
+                        value = expr_from_string(
+                            f"tl.load({value.id} + offset, mask)",
+                            offset=value_indexing.index_expr,
+                            mask=value_indexing.mask_expr,
+                        )
+                    break
+        
         return expr_from_string(
             f"tl.store({name} + offset, value, mask)",
             value=value,
@@ -511,7 +530,14 @@ def compute_shape(
                 output_size.extend(k.size())
             else:
                 raise exc.InvalidIndexingType(k)
-        assert len(input_size) == 0, "invalid subscript"
+        # For partial indexing, append remaining dimensions to output
+        while input_size:
+            size = input_size.popleft()
+            if size != 1:
+                rdim = env.allocate_reduction_dimension(size)
+                output_size.append(rdim.var)
+            else:
+                output_size.append(1)
         return output_size
 
     @staticmethod
@@ -648,6 +674,22 @@ def create(
                         )
             else:
                 raise exc.InvalidIndexingType(type(k))
+        
+        # Handle remaining dimensions for partial indexing
+        while len(index_values) < fake_value.ndim:
+            expand = tile_strategy.expand_str(output_size, output_idx)
+            size = fake_value.size(len(index_values))
+            if size != 1:
+                rdim = env.allocate_reduction_dimension(size)
+                block_idx = rdim.block_id
+                index_var = state.codegen.index_var(block_idx)
+                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}")
+            output_idx += 1
+            
         assert len(output_size) == output_idx
         assert len(index_values) == fake_value.ndim
         index_expr = []

helion/_compiler/type_propagation.py

@@ -487,12 +487,17 @@ def _device_indexing_size(self, key: TypeInfo) -> list[int | torch.SymInt]:
                 raise exc.OverpackedTile(k)
             else:
                 raise exc.InvalidIndexingType(k)
-        if inputs_consumed != self.fake_value.ndim:
-            raise exc.RankMismatch(
-                self.fake_value.ndim,
-                inputs_consumed,
-                f"tensor shape: {tuple(self.fake_value.shape)}",
-            )
+        # Handle partial indexing - add remaining dimensions to output
+        if inputs_consumed < self.fake_value.ndim:
+            for i in range(inputs_consumed, self.fake_value.ndim):
+                size = self.fake_value.size(i)
+                if self.origin.is_device():
+                    output_sizes.append(size)
+                elif size != 1:
+                    rdim = env.allocate_reduction_dimension(size)
+                    output_sizes.append(rdim.var)
+                else:
+                    output_sizes.append(1)
         return output_sizes
 
     def propagate_setitem(

helion/language/_tracing_ops.py

@@ -68,7 +68,11 @@ def _host_tensor(debug_name: str) -> torch.Tensor:
 
 @_decorators.codegen(_host_tensor)
 def _(state: CodegenState) -> ast.AST:
-    return expr_from_string("_host_tensor")  # should be unused
+    # Get the tensor from the FX node metadata
+    tensor = state.fx_node.meta["val"]  # pyright: ignore[reportOptionalMemberAccess]
+    # Get the tensor argument name from the device function
+    tensor_arg = state.device_function.tensor_arg(tensor)
+    return expr_from_string(tensor_arg.name)
 
 
 @has_side_effect

test/test_indexing.py

@@ -808,9 +808,6 @@ def kernel(
         torch.testing.assert_close(src_result, expected_src)
         torch.testing.assert_close(dst_result, expected_dst)
 
-    @skipIfNormalMode(
-        "RankMismatch: Expected ndim=2, but got ndim=1 - tensor value assignment shape mismatch"
-    )
     def test_tensor_value(self):
         """Test both setter from tensor value and getter for [i]"""