Add hl.reduce

stack-info: PR: #240, branch: jansel/stack/79

Author: jansel

helion/_compiler/device_ir.py

@@ -740,6 +740,9 @@ def visit_Assign(self, node: ast.Assign) -> None:
             for t, v in zip(target.elts, value, strict=True):
                 if isinstance(t, ast.Name):
                     self._assign(t, v)
+                elif isinstance(t, ast.Subscript):
+                    # Handle subscript targets in tuple unpacking (e.g., a[i], b[j] = tuple)
+                    self._assign_subscript(t, v)
                 else:
                     raise exc.InvalidAssignment
             return None
@@ -757,6 +760,23 @@ def visit_Assign(self, node: ast.Assign) -> None:
         target_origin = target.value._type_info.origin
         assert target_origin.is_host()
         val = self.visit(node.value)
+        self._assign_subscript(target, val)
+
+    def _assign_subscript(self, target: ast.Subscript, val: object) -> None:
+        """Helper method to assign a value to a subscript target."""
+        assert isinstance(target, ExtendedAST)
+        lhs_type = target._type_info
+
+        # Validate that we're assigning to a tensor subscript
+        from .type_propagation import TensorType
+
+        if not isinstance(lhs_type, TensorType):
+            raise exc.NonTensorSubscriptAssign(lhs_type, type(val))
+
+        assert isinstance(target.value, ExtendedAST)
+        target_origin = target.value._type_info.origin
+        assert target_origin.is_host()
+
         return hl.store(
             self.visit(target.value), self._subscript_slice_proxy(target.slice), val
         )

helion/_compiler/roll_reduction.py

@@ -15,6 +15,7 @@
 from helion.language._tracing_ops import _host_tensor
 from helion.language._tracing_ops import _if
 from helion.language.memory_ops import store
+from helion.language.reduce_ops import _reduce
 
 if TYPE_CHECKING:
     from helion._compiler.compile_environment import BlockSizeInfo
@@ -69,6 +70,12 @@ def should_go_in_inner_graph(self, node: torch.fx.Node) -> bool:
             return False
         assert node.op == "call_function", f"Unsupported node type {node.op}"
 
+        if node.target is _reduce:
+            # TODO(jansel): support rolling user-defined reductions
+            raise NotImplementedError(
+                "hl._reduce operations are not compatible with reduction rolling"
+            )
+
         if node.target in (_for_loop, _if):
             if node.target is _for_loop:
                 graph_id, *_ = node.args

helion/language/__init__.py

@@ -11,6 +11,7 @@
 from .memory_ops import atomic_add as atomic_add
 from .memory_ops import load as load
 from .memory_ops import store as store
+from .reduce_ops import reduce as reduce
 from .scan_ops import associative_scan as associative_scan
 from .scan_ops import cumprod as cumprod
 from .scan_ops import cumsum as cumsum