wip

Author: yf225

helion/_compiler/device_function.py

@@ -282,8 +282,8 @@ def block_size_var(self, block_id: int) -> str | None:
 
             var_name = self.new_var(f"_BLOCK_SIZE_{block_id}")
             self.block_size_var_cache[key] = var_name
-            host_expr = HostFunction.current().literal_expr(block_value)
-            if self.constexpr_arg(var_name, host_expr):
+            if self.constexpr_arg(var_name, block_value):
+                host_expr = self._constexpr_args[var_name].host_str()
                 self.codegen.host_statements.append(
                     statement_from_string(f"{var_name} = {host_expr}")
                 )
@@ -484,14 +484,42 @@ def expr_arg(self, sym: sympy.Expr, origin: Origin) -> SymbolArgument:
             self._expr_args[sym] = arg
         return self._expr_args[sym]
 
-    def constexpr_arg(self, name: str, host_str: str | None = None) -> bool:
+    def constexpr_arg(self, name: str, value: object | None = None) -> bool:
         """Create a constexpr argument, returns True if created, False if already exists."""
         if name in self._constexpr_args:
             return False
-        self._constexpr_args[name] = rv = ConstExprArg(name, host_str or name)
+        host_str = name if value is None else self._format_constexpr_value(value)
+        self._constexpr_args[name] = rv = ConstExprArg(name, host_str)
         self.arguments.append(rv)
         return True
 
+    def _format_constexpr_value(self, value: object) -> str:
+        if isinstance(value, str):
+            return value
+        if isinstance(value, (int, float, bool)):
+            return repr(value)
+
+        # Extract sympy expression from torch symbolic types
+        if isinstance(value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
+            value = value._sympy_()
+
+        # Handle sympy expressions (sanitize by replacing triton_helpers functions)
+        if isinstance(value, sympy.Expr):
+            expr = value.replace(
+                lambda node: isinstance(node, sympy.Function)
+                and getattr(node.func, "__name__", "")
+                == "triton_helpers.div_floor_integer",
+                lambda node: sympy.floor(node.args[0] / node.args[1]),  # pyright: ignore[reportAttributeAccessIssue]
+            ).replace(
+                lambda node: isinstance(node, sympy.Function)
+                and getattr(node.func, "__name__", "")
+                == "triton_helpers.remainder_integer",
+                lambda node: sympy.Mod(node.args[0], node.args[1]),  # pyright: ignore[reportAttributeAccessIssue]
+            )
+            return HostFunction.current().sympy_expr(expr)
+
+        return HostFunction.current().literal_expr(value)
+
     def _tensor_property(
         self,
         prop_cls: type[_P],

helion/_compiler/inductor_lowering.py

@@ -1241,8 +1241,7 @@ def _create_named_result(self, node: Node, result: ast.expr) -> str:
         ):
             # This expression is used in tl.arange, make it a constexpr
             name = self.cg.device_function.new_var(node.name)
-            host_expr = self.cg.device_function.sympy_expr(val._sympy_())
-            self.cg.device_function.constexpr_arg(name, host_expr)
+            self.cg.device_function.constexpr_arg(name, val._sympy_())
             return name
 
         # If the lowering produced a named value that is already defined elsewhere

helion/_compiler/tile_strategy.py

@@ -244,11 +244,10 @@ def _setup_block_size_constexpr(
         self, state: CodegenState, block_size_var: str, block_size: SymIntLike
     ) -> None:
         """Helper to setup constexpr block size variable on host."""
-        if state.device_function.constexpr_arg(block_size_var):
+        if state.device_function.constexpr_arg(block_size_var, block_size):
+            host_expr = state.device_function._constexpr_args[block_size_var].host_str()
             state.codegen.host_statements.append(
-                statement_from_string(
-                    f"{block_size_var} = {HostFunction.current().literal_expr(block_size)}"
-                )
+                statement_from_string(f"{block_size_var} = {host_expr}")
             )