Support scalar FP8 conversion in CuTe

helion/_compiler/backend.py

@@ -211,7 +211,12 @@ def program_id_expr(self, dim: int, *, index_dtype: str) -> str:
     def cdiv_expr(self, numel: str, block_size: str, *, is_device: bool) -> str:
         return f"(({numel}) + ({block_size}) - 1) // ({block_size})"
 
-    def cast_expr(self, expr_str: str, dtype_str: str) -> str:
+    def cast_expr(
+        self,
+        expr_str: str,
+        dtype_str: str,
+        source_dtype_str: str | None = None,
+    ) -> str:
         """Generate a backend-specific type cast expression."""
         raise exc.BackendUnsupported(self.name, "cast")
 
@@ -638,9 +643,18 @@ def argreduce_loop_update_statements(
     def inductor_op_overrides(self) -> InductorOpOverrides:
         raise exc.BackendUnsupported(self.name, "Inductor OpOverrides")
 
-    def cast_ast(self, x: ast.AST, target_dtype: torch.dtype) -> ast.AST:
+    def cast_ast(
+        self,
+        x: ast.AST,
+        target_dtype: torch.dtype,
+        source_dtype: torch.dtype | None = None,
+    ) -> ast.AST:
         return expr_from_string(
-            self.cast_expr("{x}", self.dtype_str(target_dtype)),
+            self.cast_expr(
+                "{x}",
+                self.dtype_str(target_dtype),
+                self.dtype_str(source_dtype) if source_dtype is not None else None,
+            ),
             x=x,
         )
 

helion/_compiler/cute/backend.py

@@ -1038,7 +1038,9 @@ def library_imports(self) -> dict[str, str]:
             "_cute_issue_clc_query_nomulticast": "from helion._compiler.cute.clc_helpers import issue_clc_query_nomulticast as _cute_issue_clc_query_nomulticast",
             "_cute_inline_asm_elementwise": "from helion._compiler.cute.inline_asm_helpers import inline_asm_elementwise as _cute_inline_asm_elementwise",
             "_cute_fp8e4m3fn_to_float32": "from helion._compiler.cute.quantized_helpers import fp8e4m3fn_to_float32 as _cute_fp8e4m3fn_to_float32",
+            "_cute_fp8e4m3fn_to_storage": "from helion._compiler.cute.quantized_helpers import fp8e4m3fn_to_storage as _cute_fp8e4m3fn_to_storage",
             "_cute_fp8e4m3fn_x2_to_float32": "from helion._compiler.cute.quantized_helpers import fp8e4m3fn_x2_to_float32 as _cute_fp8e4m3fn_x2_to_float32",
+            "_cute_float32_to_fp8e4m3fn": "from helion._compiler.cute.quantized_helpers import float32_to_fp8e4m3fn as _cute_float32_to_fp8e4m3fn",
             "_cute_float4_e2m1fn_x2_to_float32": "from helion._compiler.cute.quantized_helpers import float4_e2m1fn_x2_to_float32 as _cute_float4_e2m1fn_x2_to_float32",
             "_cute_grid_barrier": "from helion._compiler.cute.grid_barrier import grid_barrier as _cute_grid_barrier",
             "_cute_atomic_max_float32": "from helion._compiler.cute.atomic_helpers import atomic_max_float32 as _cute_atomic_max_float32",
@@ -1072,10 +1074,34 @@ def where(
 
         return HelionCuteDSLOpOverrides()
 
-    def cast_expr(self, expr_str: str, dtype_str: str) -> str:
+    def cast_expr(
+        self,
+        expr_str: str,
+        dtype_str: str,
+        source_dtype_str: str | None = None,
+    ) -> str:
+        fp8_dtype = "cutlass.Float8E4M3FN"
+        if (
+            dtype_str == fp8_dtype
+            and source_dtype_str is not None
+            and source_dtype_str not in ("cutlass.Float32", fp8_dtype)
+        ):
+            expr_str = self.cast_expr(
+                expr_str,
+                "cutlass.Float32",
+                source_dtype_str,
+            )
+            source_dtype_str = "cutlass.Float32"
+        if source_dtype_str == "cutlass.Float32" and dtype_str == fp8_dtype:
+            return f"_cute_float32_to_fp8e4m3fn({expr_str})"
         return f"{dtype_str}({expr_str})"
 
-    def cast_ast(self, x: ast.AST, target_dtype: torch.dtype) -> ast.AST:
+    def cast_ast(
+        self,
+        x: ast.AST,
+        target_dtype: torch.dtype,
+        source_dtype: torch.dtype | None = None,
+    ) -> ast.AST:
         from ..device_function import DeviceFunction
         from ..device_function import NoCurrentFunction
 
@@ -1090,7 +1116,7 @@ def cast_ast(self, x: ast.AST, target_dtype: torch.dtype) -> ast.AST:
                     is not None
                 ):
                     return x
-        return super().cast_ast(x, target_dtype)
+        return super().cast_ast(x, target_dtype, source_dtype)
 
     def grid_barrier_stmt(self, sem_arg: str) -> str | None:
         # ``sem_arg`` is a TensorArg that arrives as a ``cute.Tensor``; its

helion/_compiler/cute/memory_ops.py

@@ -156,8 +156,13 @@ def _cute_scalar_storage_dtype(dtype: torch.dtype) -> str:
 
 
 def _cute_scalar_store_expr(
-    tensor_name: str, index_exprs: list[str], value: str
+    tensor_name: str,
+    index_exprs: list[str],
+    value: str,
+    dtype: torch.dtype,
 ) -> str:
+    if dtype is torch.float8_e4m3fn:
+        value = f"_cute_fp8e4m3fn_to_storage({value})"
     if "None" in index_exprs:
         return f"{tensor_name}.__setitem__({_cute_index_tuple(index_exprs)}, {value})"
     return f"{_cute_scalar_pointer_expr(tensor_name, index_exprs)}.store({value})"
@@ -427,6 +432,7 @@ def _codegen_cute_store_stack_load(
             tensor_name,
             target_indices,
             f"({stack_ptr_expr}).load()",
+            tensor.dtype,
         )
         mask_expr = _cute_combined_mask(state, [*subscript], extra_mask, tensor=tensor)
         if mask_expr is None:
@@ -487,7 +493,9 @@ def _codegen_cute_store_stack_load(
         value=value,
     )
     store_expr = expr_from_string(
-        _cute_scalar_store_expr(tensor_name, rewritten_index_exprs, "{value}"),
+        _cute_scalar_store_expr(
+            tensor_name, rewritten_index_exprs, "{value}", tensor.dtype
+        ),
         value=value,
     )
     mask_expr = _cute_combined_mask(state, [*subscript], extra_mask, tensor=tensor)
@@ -755,7 +763,9 @@ def _codegen_cute_affine_reshape_store(
     )
     col_index = f"{index_dtype}({n_global})"
     tensor_name = state.device_function.tensor_arg(tensor).name
-    store_expr = _cute_scalar_store_expr(tensor_name, [row_index, col_index], value_var)
+    store_expr = _cute_scalar_store_expr(
+        tensor_name, [row_index, col_index], value_var, tensor.dtype
+    )
 
     store_stmt: ast.stmt = create(ast.Expr, value=expr_from_string(store_expr))
     mask_parts = [
@@ -1353,7 +1363,7 @@ def _codegen_cute_store_expand_broadcast_tile(
         value=value,
     )
     store_expr = expr_from_string(
-        _cute_scalar_store_expr(tensor_name, index_exprs, "{value}"),
+        _cute_scalar_store_expr(tensor_name, index_exprs, "{value}", tensor.dtype),
         value=value,
     )
 
@@ -1713,7 +1723,9 @@ def _(state: CodegenState) -> ast.AST:
     if isinstance(topk_lane_expr, str) and isinstance(topk_k, int):
         index_exprs[-1] = topk_lane_expr
     store_uses_pointer = "None" not in index_exprs
-    store_expr = _cute_scalar_store_expr(tensor_name, index_exprs, "{value}")
+    store_expr = _cute_scalar_store_expr(
+        tensor_name, index_exprs, "{value}", tensor.dtype
+    )
     assign_expr = expr_from_string(store_expr, value=value)
 
     mask_expr = _cute_combined_mask(state, subscript, extra_mask, tensor=tensor)

helion/_compiler/cute/quantized_helpers.py

@@ -4,27 +4,37 @@
 from typing import Any
 from typing import cast
 
+from cutlass import Float8E4M3FN
 from cutlass import Float32
 from cutlass import Int8
 from cutlass import Int16
+from cutlass._mlir.dialects import arith
 from cutlass._mlir.dialects import llvm
 from cutlass.cutlass_dsl import dsl_user_op
 
 if TYPE_CHECKING:
     from cutlass._mlir import ir
 
 
-def _as_i16(
+def _as_ir_value(
     value: object,
     *,
     loc: ir.Location | None,
     ip: ir.InsertionPoint | None,
 ) -> ir.Value:
     raw_value = getattr(value, "value", None)
     if hasattr(raw_value, "type"):
-        ir_value = cast("Any", raw_value)
-    else:
-        ir_value = cast("Any", value).ir_value(loc=loc, ip=ip)
+        return cast("Any", raw_value)
+    return cast("Any", value).ir_value(loc=loc, ip=ip)
+
+
+def _as_i16(
+    value: object,
+    *,
+    loc: ir.Location | None,
+    ip: ir.InsertionPoint | None,
+) -> ir.Value:
+    ir_value = _as_ir_value(value, loc=loc, ip=ip)
     ir_type = str(cast("Any", ir_value).type)
     if ir_type == "i8":
         return llvm.zext(Int16.mlir_type, ir_value, loc=loc, ip=ip)
@@ -36,6 +46,56 @@ def _as_i16(
     raise TypeError(f"unsupported quantized scalar type: {ir_type}")
 
 
+@dsl_user_op
+def float32_to_fp8e4m3fn(
+    value: Any,  # noqa: ANN401
+    *,
+    loc: ir.Location | None = None,
+    ip: ir.InsertionPoint | None = None,
+) -> Float8E4M3FN:
+    """Narrow a scalar through PTX's packed E4M3 conversion instruction."""
+    if isinstance(value, Float8E4M3FN):
+        return value
+
+    value_f32 = Float32(value).ir_value(loc=loc, ip=ip)
+    packed = llvm.inline_asm(
+        Int16.mlir_type,
+        [value_f32],
+        """
+        {
+          .reg .f32 zero;
+          mov.f32 zero, 0f00000000;
+          cvt.rn.satfinite.e4m3x2.f32 $0, zero, $1;
+        }
+        """,
+        "=h,f",
+        has_side_effects=False,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    low_byte = arith.trunci(Int8.mlir_type, packed, loc=loc, ip=ip)
+    return Float8E4M3FN(llvm.bitcast(Float8E4M3FN.mlir_type, low_byte, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def fp8e4m3fn_to_storage(
+    value: object,
+    *,
+    loc: ir.Location | None = None,
+    ip: ir.InsertionPoint | None = None,
+) -> Int8:
+    """Return the raw storage byte for a scalar E4M3 register value."""
+    ir_value = _as_ir_value(value, loc=loc, ip=ip)
+    ir_type = str(cast("Any", ir_value).type)
+    if ir_type == "i8":
+        return Int8(ir_value)
+    if ir_type.startswith("f8"):
+        return Int8(llvm.bitcast(Int8.mlir_type, ir_value, loc=loc, ip=ip))
+    raise TypeError(f"unsupported fp8 storage value type: {ir_type}")
+
+
 @dsl_user_op
 def fp8e4m3fn_to_float32(
     value: object,

helion/_compiler/dtype_utils.py

@@ -4,19 +4,20 @@
 
 import torch
 
-from .ast_extension import expr_from_string
 from .compile_environment import CompileEnvironment
 
 if TYPE_CHECKING:
     import ast
 
 
-def cast_ast(x: ast.AST, dtype: torch.dtype) -> ast.AST:
+def cast_ast(
+    x: ast.AST,
+    dtype: torch.dtype,
+    source_dtype: torch.dtype | None = None,
+) -> ast.AST:
     """Return an AST that casts expression `x` to the backend dtype string."""
     env = CompileEnvironment.current()
-    dtype_str = env.backend.dtype_str(dtype)
-    cast_str = env.backend.cast_expr("{x}", dtype_str)
-    return expr_from_string(cast_str, x=x)
+    return env.backend.cast_ast(x, dtype, source_dtype)
 
 
 def promote_and_cast_pair(
@@ -32,6 +33,6 @@ def promote_and_cast_pair(
     """
 
     common = torch.promote_types(lhs_dtype, rhs_dtype)
-    lhs_out = lhs if lhs_dtype == common else cast_ast(lhs, common)
-    rhs_out = rhs if rhs_dtype == common else cast_ast(rhs, common)
+    lhs_out = lhs if lhs_dtype == common else cast_ast(lhs, common, lhs_dtype)
+    rhs_out = rhs if rhs_dtype == common else cast_ast(rhs, common, rhs_dtype)
     return lhs_out, rhs_out, common

helion/_compiler/inductor_lowering.py

@@ -1072,9 +1072,14 @@ def __init__(
         self.cg = cg
         self.input_name_lookup = input_name_lookup
 
-    def _cast_ast(self, x: ast.AST, target_dtype: torch.dtype) -> ast.AST:
+    def _cast_ast(
+        self,
+        x: ast.AST,
+        target_dtype: torch.dtype,
+        source_dtype: torch.dtype | None = None,
+    ) -> ast.AST:
         backend = CompileEnvironment.current().backend
-        return backend.cast_ast(x, target_dtype)
+        return backend.cast_ast(x, target_dtype, source_dtype)
 
     def _to_ast(self, x: object) -> ast.AST:
         if isinstance(x, ast.AST):
@@ -1094,7 +1099,12 @@ def _expected_tensor_dtype(self) -> torch.dtype | None:
             return val.dtype
         return None
 
-    def _create_cast_expr(self, x: object, target_dtype: torch.dtype) -> ast.AST:
+    def _create_cast_expr(
+        self,
+        x: object,
+        target_dtype: torch.dtype,
+        source_dtype: torch.dtype | None = None,
+    ) -> ast.AST:
         """Create a backend cast expression from AST or string input.
 
         Args:
@@ -1105,7 +1115,7 @@ def _create_cast_expr(self, x: object, target_dtype: torch.dtype) -> ast.AST:
             AST expression for the cast operation
         """
         x_ast = self._to_ast(x)
-        return self._cast_ast(x_ast, target_dtype)
+        return self._cast_ast(x_ast, target_dtype, source_dtype)
 
     def _maybe_cast_to_expected_dtype(self, expr: ast.AST) -> ast.AST:
         """Cast expression to expected dtype if needed.
@@ -1156,7 +1166,7 @@ def to_dtype(
                 x=self._to_ast(x),
             )
             return self._lift(cast_expr)
-        cast_expr = self._create_cast_expr(x, dtype)
+        cast_expr = self._create_cast_expr(x, dtype, src_dtype)
         return self._lift(cast_expr)
 
     def sigmoid(self, x: object) -> str:  # type: ignore[override]

helion/_compiler/metal/backend.py

@@ -103,7 +103,12 @@ def index_type_str(self, index_dtype: torch.dtype) -> str:
     def inline_constexpr(self, name: str, value: str) -> str:
         return f"{name} = {value}"
 
-    def cast_expr(self, expr_str: str, dtype_str: str) -> str:
+    def cast_expr(
+        self,
+        expr_str: str,
+        dtype_str: str,
+        source_dtype_str: str | None = None,
+    ) -> str:
         return f"static_cast<{dtype_str}>({expr_str})"
 
     def lane_index_expr(

helion/_compiler/pallas/backend.py

@@ -198,7 +198,12 @@ def supports_config_key(self, key: str) -> bool:
     def program_id_expr(self, dim: int, *, index_dtype: str) -> str:
         return f"pl.program_id({dim})"
 
-    def cast_expr(self, expr_str: str, dtype_str: str) -> str:
+    def cast_expr(
+        self,
+        expr_str: str,
+        dtype_str: str,
+        source_dtype_str: str | None = None,
+    ) -> str:
         return f"lax.convert_element_type({expr_str}, {dtype_str})"
 
     @property
@@ -240,7 +245,12 @@ def inductor_op_overrides(self) -> InductorOpOverrides:
 
         return PallasKernelOverrides()
 
-    def cast_ast(self, x: ast.AST, target_dtype: torch.dtype) -> ast.AST:
+    def cast_ast(
+        self,
+        x: ast.AST,
+        target_dtype: torch.dtype,
+        source_dtype: torch.dtype | None = None,
+    ) -> ast.AST:
         return expr_from_string(
             f"lax.convert_element_type({{x}}, {self.dtype_str(target_dtype)})", x=x
         )

helion/_compiler/triton/backend.py

@@ -234,7 +234,12 @@ def acc_type(self, dtype: torch.dtype) -> str:
 
         return triton_acc_type(dtype)
 
-    def cast_expr(self, expr_str: str, dtype_str: str) -> str:
+    def cast_expr(
+        self,
+        expr_str: str,
+        dtype_str: str,
+        source_dtype_str: str | None = None,
+    ) -> str:
         return f"tl.cast({expr_str}, {dtype_str})"
 
     def arange_expr(