diff --git a/helion/_compiler/backend.py b/helion/_compiler/backend.py index 144145462..4728f97c3 100644 --- a/helion/_compiler/backend.py +++ b/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, ) diff --git a/helion/_compiler/cute/backend.py b/helion/_compiler/cute/backend.py index 496236a1f..7879117cf 100644 --- a/helion/_compiler/cute/backend.py +++ b/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 diff --git a/helion/_compiler/cute/memory_ops.py b/helion/_compiler/cute/memory_ops.py index 0e055b048..c5f386d25 100644 --- a/helion/_compiler/cute/memory_ops.py +++ b/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) diff --git a/helion/_compiler/cute/quantized_helpers.py b/helion/_compiler/cute/quantized_helpers.py index e60abf948..e6dbffce6 100644 --- a/helion/_compiler/cute/quantized_helpers.py +++ b/helion/_compiler/cute/quantized_helpers.py @@ -4,9 +4,11 @@ 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 @@ -14,7 +16,7 @@ from cutlass._mlir import ir -def _as_i16( +def _as_ir_value( value: object, *, loc: ir.Location | None, @@ -22,9 +24,17 @@ def _as_i16( ) -> 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, diff --git a/helion/_compiler/dtype_utils.py b/helion/_compiler/dtype_utils.py index b6f879494..ec3855f80 100644 --- a/helion/_compiler/dtype_utils.py +++ b/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 diff --git a/helion/_compiler/inductor_lowering.py b/helion/_compiler/inductor_lowering.py index 8fadc985a..630434ff4 100644 --- a/helion/_compiler/inductor_lowering.py +++ b/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] diff --git a/helion/_compiler/metal/backend.py b/helion/_compiler/metal/backend.py index 4ed786ba8..369f810f5 100644 --- a/helion/_compiler/metal/backend.py +++ b/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( diff --git a/helion/_compiler/pallas/backend.py b/helion/_compiler/pallas/backend.py index 741ecfe6c..09a51f6b8 100644 --- a/helion/_compiler/pallas/backend.py +++ b/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 ) diff --git a/helion/_compiler/triton/backend.py b/helion/_compiler/triton/backend.py index fa930ab85..fd201c67f 100644 --- a/helion/_compiler/triton/backend.py +++ b/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( diff --git a/test/test_cute_backend.py b/test/test_cute_backend.py index 85905f948..98aeab1ee 100644 --- a/test/test_cute_backend.py +++ b/test/test_cute_backend.py @@ -20,6 +20,7 @@ from helion._testing import TestCase from helion._testing import code_and_output from helion._testing import onlyBackends +from helion._testing import skipIfCudaCapabilityLessThan from helion.exc import BackendUnsupported from helion.exc import CuteBackendUnavailable import helion.language as hl @@ -121,6 +122,14 @@ def cute_relu(x: torch.Tensor) -> torch.Tensor: return out +@helion.kernel(backend="cute") +def cute_cast_fp8(x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x, dtype=torch.float8_e4m3fn) + for tile in hl.tile(out.size()): + out[tile] = x[tile].to(torch.float8_e4m3fn) + return out + + @helion.kernel(backend="cute") def cute_sin(x: torch.Tensor) -> torch.Tensor: out = torch.empty_like(x) @@ -4173,6 +4182,37 @@ def test_pointwise_relu(self) -> None: (x,) = args torch.testing.assert_close(out, torch.relu(x)) + @skipIfCudaCapabilityLessThan((9, 0), reason="FP8 requires CUDA capability >= 9.0") + def test_pointwise_cast_fp8(self) -> None: + x = torch.linspace(-10, 10, 32, device=DEVICE) + code, out = code_and_output( + cute_cast_fp8, + (x,), + block_size=32, + num_threads=[32], + ) + self.assertEqual(out.dtype, torch.float8_e4m3fn) + torch.testing.assert_close( + out.float(), + x.to(torch.float8_e4m3fn).float(), + ) + self.assertIn("_cute_float32_to_fp8e4m3fn", code) + self.assertIn("_cute_fp8e4m3fn_to_storage", code) + + x_bf16 = x.to(torch.bfloat16) + bf16_code, bf16_out = code_and_output( + cute_cast_fp8, + (x_bf16,), + block_size=32, + num_threads=[32], + ) + self.assertEqual(bf16_out.dtype, torch.float8_e4m3fn) + torch.testing.assert_close( + bf16_out.float(), + x_bf16.to(torch.float8_e4m3fn).float(), + ) + self.assertIn("_cute_float32_to_fp8e4m3fn(cutlass.Float32(load))", bf16_code) + def test_pointwise_sin(self) -> None: args = (torch.randn(65, 23, device=DEVICE, dtype=torch.float32),) code, out = code_and_output(cute_sin, args) diff --git a/test/test_indexing.py b/test/test_indexing.py index efc359870..94bd01729 100644 --- a/test/test_indexing.py +++ b/test/test_indexing.py @@ -1358,7 +1358,7 @@ def size1_reshape_kernel( torch.testing.assert_close(out2, x2) @skipIfCudaCapabilityLessThan((9, 0), reason="FP8 requires CUDA capability >= 9.0") - @skipIfCute("CuTe does not support scalar float32 to float8 conversion") + @skipIfCute("CuTe does not support this reshape/reduction layout") @skipIfRefEager("Test validates compiler block-size constraints") def test_tile_reshape_for_grouped_quantization(self): group_size = 128