Add support for listcomp (#412)

Author: jansel

helion/_compiler/device_ir.py

@@ -475,7 +475,7 @@ def disable_tracing() -> Iterator[proxy_tensor.PythonKeyTracer]:
 
     @staticmethod
     def should_become_arg(value: object) -> bool:
-        if isinstance(value, (Tile, torch.SymInt)):
+        if isinstance(value, (Tile, int, float, bool, type(None), torch.SymInt)):
             return False
         if isinstance(value, torch.Tensor):
             if (
@@ -502,20 +502,54 @@ def _extract_tile_begin_end(self, for_node: ast.For) -> tuple[object, object]:
             end = self.visit(args[1])
         return begin, end
 
+    def _handle_sequence_unrolling(
+        self,
+        sequence_iter: ast.AST,
+        target: ast.AST,
+        element_processor: Callable[[], object | None],
+        preserve_scope: bool = False,
+    ) -> list[object]:
+        """Common logic for unrolling sequences in both loops and comprehensions."""
+        # Get the sequence of values to iterate over
+        sequence_value = self.visit(sequence_iter)
+        assert isinstance(sequence_value, (tuple, list)), (
+            f"Expected tuple or list, got {type(sequence_value)}"
+        )
+
+        results = []
+        for element_value in sequence_value:
+            if preserve_scope:
+                # For loops: don't create new scope, allow state to persist
+                self._assign(target, element_value)
+                result = element_processor()
+                if result is not None:
+                    results.append(result)
+            else:
+                # For comprehensions: create isolated scope for each iteration
+                old_scope = self.scope.copy()
+                try:
+                    self._assign(target, element_value)
+                    result = element_processor()
+                    if result is not None:
+                        results.append(result)
+                finally:
+                    self.scope = old_scope
+
+        return results
+
     def _handle_tuple_unrolling(
         self,
         node: ast.For,
     ) -> None:
         """Handle unrolling of loops that iterate over tuples of tensors."""
-        # Get the sequence of tensors to iterate over
-        sequence_value = self.visit(node.iter)
-        assert isinstance(sequence_value, (tuple, list)), (
-            f"Expected tuple or list, got {type(sequence_value)}"
-        )
-        # Unroll the loop by executing the body for each tensor in the sequence
-        for tensor_value in sequence_value:
-            self._assign(node.target, tensor_value)
+
+        def execute_body() -> None:
             self._body(node.body)
+            return None  # No result to collect for loops
+
+        self._handle_sequence_unrolling(
+            node.iter, node.target, execute_body, preserve_scope=True
+        )
 
     def visit_For(self, node: ast.For) -> None:
         assert isinstance(node, ExtendedAST)
@@ -528,6 +562,15 @@ def visit_For(self, node: ast.For) -> None:
             self._handle_tuple_unrolling(node)
             return
 
+        # Special handling for variables that might contain sequences from list comprehensions
+        if isinstance(node.iter, ast.Name) and node.iter.id in self.scope:
+            scope_value = self.scope[node.iter.id]
+            if isinstance(scope_value, (tuple, list)):
+                # This is a sequence in the scope, we should try to unroll it
+                # even if the type info doesn't indicate it's a SequenceType
+                self._handle_tuple_unrolling(node)
+                return
+
         if not isinstance(iter_type, IterType):
             raise exc.InvalidDeviceForLoop(iter_type)
         inner_type: TypeInfo = iter_type.inner
@@ -724,6 +767,50 @@ def visit_Tuple(self, node: ast.Tuple) -> tuple[object, ...]:
     def visit_List(self, node: ast.List) -> list[object]:
         return [self.visit(x) for x in node.elts]
 
+    def visit_ListComp(self, node: ast.ListComp) -> tuple[object, ...]:
+        """Handle list comprehension unrolling similar to tuple unrolling."""
+        assert isinstance(node, ExtendedAST)
+
+        # Only handle simple cases with single generator and no if conditions
+        if len(node.generators) != 1 or node.generators[0].ifs:
+            raise exc.StatementNotSupported(
+                "Complex list comprehensions are not supported"
+            )
+
+        generator = node.generators[0]
+        assert isinstance(generator.iter, ExtendedAST)
+        iter_type = generator.iter._type_info
+
+        # Check if we're iterating over a sequence (similar to tuple unrolling)
+        if isinstance(iter_type, SequenceType):
+            return self._handle_listcomp_unrolling(node)
+
+        # For non-sequence iterables, we could extend this later
+        raise exc.StatementNotSupported(
+            "List comprehensions over non-sequence types are not supported"
+        )
+
+    def _handle_listcomp_unrolling(self, node: ast.ListComp) -> tuple[object, ...]:
+        """Handle unrolling of list comprehensions over sequences."""
+        generator = node.generators[0]
+
+        def evaluate_expression() -> object:
+            # Evaluate the comprehension expression
+            result = self.visit(node.elt)
+            # If the result is a SymInt that can be evaluated to a concrete value, do so
+            if isinstance(result, torch.SymInt):
+                try:
+                    return int(result)
+                except (ValueError, TypeError):
+                    return result
+            return result
+
+        results = self._handle_sequence_unrolling(
+            generator.iter, generator.target, evaluate_expression, preserve_scope=False
+        )
+        # Return as tuple to match the expected type for tuple unrolling
+        return tuple(results)
+
     def visit_Slice(self, node: ast.Slice) -> slice:
         if node.lower is None:
             lower = None

helion/_compiler/inductor_lowering.py

@@ -969,6 +969,30 @@ def _default(
         result_str = _unpack_opsvalue(
             getattr(self.parent_handler, name)(*args, **kwargs)
         )
+
+        return self.cg.lift(expr_from_string(result_str)).id
+
+    def to_dtype(
+        self,
+        x: object,
+        dtype: torch.dtype,
+        src_dtype: torch.dtype | None = None,
+        use_compute_types: bool = True,
+    ) -> str:
+        """Override to_dtype to use tl.cast for scalar values from GetItemOrigin."""
+        x_str = str(x)
+
+        # Use tl.cast for scalar values (typically from GetItemOrigin)
+        # These are plain scalars that should use tl.cast instead of .to()
+        if "_item_" in x_str:
+            return self.cg.lift(
+                expr_from_string(f"tl.cast({x_str}, {triton_type(dtype)})")
+            ).id
+
+        # Fall back to the default behavior for other cases
+        result_str = _unpack_opsvalue(
+            self.parent_handler.to_dtype(x, dtype, src_dtype, use_compute_types)
+        )
         return self.cg.lift(expr_from_string(result_str)).id
 
     def load(self, name: str, index: sympy.Expr) -> str:

helion/_compiler/type_propagation.py

@@ -263,6 +263,11 @@ def from_example(cls, value: object, origin: Origin) -> TypeInfo:
                     )
                 ),
             )
+        if isinstance(value, zip):
+            # Handle zip objects by converting to tuple of tuples
+            # This allows zip to work in list comprehensions
+            zipped_tuples = tuple(tuple(items) for items in value)
+            return cls.from_example(zipped_tuples, origin)
         raise exc.UnsupportedPythonType(type(value).__name__)
 
     @staticmethod
@@ -2039,8 +2044,63 @@ def _not_on_device_statement(self, node: ast.AST) -> TypeInfo:
     def _not_supported(self, node: ast.AST) -> TypeInfo:
         raise exc.StatementNotSupported(type(node).__name__)
 
+    def _evaluate_comprehension(
+        self, generator: ast.comprehension, expression: ast.AST
+    ) -> TypeInfo:
+        """Helper method to evaluate comprehension type propagation."""
+        # Visit the iterable to get its type
+        iter_type = self.visit(generator.iter)
+
+        # Get element type and evaluate expression in scope
+        self.push_scope()
+        try:
+            element_type = iter_type.propagate_iter(self.origin())
+            self._assign(generator.target, element_type)
+
+            # Process conditional filters (basic validation)
+            for if_clause in generator.ifs:
+                self.visit(if_clause)
+
+            element_result_type = self.visit(expression)
+        finally:
+            self.pop_scope()
+
+        # Try to determine exact result by unpacking iterable
+        try:
+            iterable_elements = iter_type.unpack()
+            result_elements = []
+
+            for element_type in iterable_elements:
+                self.push_scope()
+                try:
+                    self._assign(generator.target, element_type)
+                    # For now, assume all conditions pass
+                    for if_clause in generator.ifs:
+                        self.visit(if_clause)
+                    result_elements.append(self.visit(expression))
+                finally:
+                    self.pop_scope()
+
+            # If there are conditions, we can't determine exact length
+            if generator.ifs and result_elements:
+                result_elements = [result_elements[0]]
+
+            return SequenceType(self.origin(), result_elements)
+
+        except NotImplementedError:
+            # Fallback to generic list type
+            return SequenceType(self.origin(), [element_result_type])
+
+    def visit_ListComp(self, node: ast.ListComp) -> TypeInfo:
+        """Type propagation for list comprehensions."""
+        if len(node.generators) != 1:
+            raise exc.StatementNotSupported(
+                "List comprehensions with multiple generators are not supported"
+            )
+
+        return self._evaluate_comprehension(node.generators[0], node.elt)
+
     # TODO(jansel): need to implement these
-    visit_ListComp: _VisitMethod = _not_supported  # pyright: ignore[reportAssignmentType,reportIncompatibleMethodOverride]
     visit_SetComp: _VisitMethod = _not_supported  # pyright: ignore[reportAssignmentType,reportIncompatibleMethodOverride]
     visit_GeneratorExp: _VisitMethod = _not_supported  # pyright: ignore[reportAssignmentType,reportIncompatibleMethodOverride]
     visit_DictComp: _VisitMethod = _not_supported  # pyright: ignore[reportAssignmentType,reportIncompatibleMethodOverride]

test/test_misc.expected

@@ -190,7 +190,7 @@ def _tuple_literal_index_kernel_kernel(out, inp_tuple_item_0, inp_tuple_item_1,
     load_1 = tl.load(inp_tuple_item_1 + (indices_0[:, None] * inp_tuple_item_1_stride_0 + indices_1[None, :] * inp_tuple_item_1_stride_1), mask_0[:, None] & mask_1[None, :], other=0)
     v_0 = load_1.to(tl.float32)
     v_1 = load + v_0
-    v_2 = inp_tuple_item_2.to(tl.float32)
+    v_2 = tl.cast(inp_tuple_item_2, tl.float32)
     v_3 = v_1 * v_2
     tl.store(out + (indices_0[:, None] * out_stride_0 + indices_1[None, :] * out_stride_1), v_3, mask_0[:, None] & mask_1[None, :])
 
@@ -217,7 +217,7 @@ def _tuple_literal_index_kernel_kernel(out, inp_tuple_item_0, inp_tuple_item_1,
     load_1 = tl.load(tl.make_block_ptr(inp_tuple_item_1, [inp_tuple_item_1_size_0, inp_tuple_item_1_size_1], [inp_tuple_item_1_stride_0, inp_tuple_item_1_stride_1], [offset_0, offset_1], [_BLOCK_SIZE_0, _BLOCK_SIZE_1], [1, 0]), boundary_check=[0, 1], padding_option='zero')
     v_0 = load_1.to(tl.float32)
     v_1 = load + v_0
-    v_2 = inp_tuple_item_2.to(tl.float32)
+    v_2 = tl.cast(inp_tuple_item_2, tl.float32)
     v_3 = v_1 * v_2
     tl.store(tl.make_block_ptr(out, [out_size_0, out_size_1], [out_stride_0, out_stride_1], [offset_0, offset_1], [_BLOCK_SIZE_0, _BLOCK_SIZE_1], [1, 0]), v_3, boundary_check=[0, 1])
 
@@ -255,7 +255,7 @@ def _tuple_literal_index_kernel_kernel(out, inp_tuple_item_0, inp_tuple_item_1,
     load_1 = inp_tuple_item_1_desc.load([offset_0, offset_1])
     v_0 = load_1.to(tl.float32)
     v_1 = load + v_0
-    v_2 = inp_tuple_item_2.to(tl.float32)
+    v_2 = tl.cast(inp_tuple_item_2, tl.float32)
     v_3 = v_1 * v_2
     tl.store(out + (indices_0[:, None] * out_stride_0 + indices_1[None, :] * out_stride_1), v_3, mask_0[:, None] & mask_1[None, :])