"""Lower elementwise inline PTX over CuteDSL fragments. Each assembly invocation consumes `pack` adjacent fragment elements. LLVM constraints select the physical register type for every operand, so inputs are bitcast into those registers rather than numerically converted. Multiple outputs are returned as an LLVM struct and extracted before rebuilding the fragment. Partial packs are zero-padded, and outputs for padded elements are discarded. Integer register results are narrowed to the requested element width. E8M0 codes 1 through 254 map directly to the matching Float32 exponent bits; codes 0 and 255 map explicitly to `2^-127` and NaN. This lets fused arithmetic use exact Float32 values while preserving E8M0 as the stored result type. """ import functools import hashlib import math from pathlib import Path import cutlass import cutlass.cute as cute from cutlass._mlir import ir from cutlass._mlir.dialects import arith, llvm, vector from cutlass.cute.tensor import _infer_broadcast_shape from cutlass.cutlass_dsl import dsl_user_op, T @functools.cache def inline_asm_cache_key() -> str: """Include this module's source in generated kernel cache keys.""" return hashlib.sha256(Path(__file__).read_bytes()).hexdigest() # PTX register classes, as accepted by the NVPTX inline-asm backend. CONSTRAINT_TYPES = { "h": T.i16, "r": T.i32, "l": T.i64, "f": T.f32, "d": T.f64, } def split_constraints(constraints: str) -> tuple[list[str], list[str]]: """Split an LLVM constraint list into output and input register classes.""" outputs, inputs = [], [] for entry in constraints.split(","): entry = entry.strip() target = outputs if entry.startswith("=") else inputs letter = entry.lstrip("=&") if letter not in CONSTRAINT_TYPES: raise NotImplementedError( f"CuteDSL inline asm supports constraints {sorted(CONSTRAINT_TYPES)}, " f"got {entry!r}" ) target.append(letter) return outputs, inputs def bitcast_to(value: ir.Value, target: ir.Type) -> ir.Value: """Reinterpret a value's bits as an equally wide type.""" if value.type == target: return value return arith.bitcast(target, value) def narrow_integer(value: ir.Value, width: int) -> ir.Value: """Truncate an integer register to `width` bits.""" if not ir.IntegerType.isinstance(value.type): raise TypeError(f"expected an integer inline-asm result, got {value.type}") if ir.IntegerType(value.type).width > width: return arith.trunci(ir.IntegerType.get_signless(width), value) return value def decode_e8m0(value: ir.Value) -> ir.Value: """Decode one biased E8M0 bit pattern to its exact Float32 value.""" code = arith.extui(T.i32(), narrow_integer(value, 8)) bits = arith.shli(code, arith.constant(T.i32(), 23)) is_zero = arith.cmpi(arith.CmpIPredicate.eq, code, arith.constant(T.i32(), 0)) bits = arith.select(is_zero, arith.constant(T.i32(), 0x00400000), bits) is_nan = arith.cmpi(arith.CmpIPredicate.eq, code, arith.constant(T.i32(), 255)) bits = arith.select(is_nan, arith.constant(T.i32(), 0x7FC00000), bits) return arith.bitcast(T.f32(), bits) def convert_output(value: ir.Value, result_type) -> ir.Value: """Convert a register result to the requested logical element type.""" if result_type == cutlass.Float8E8M0FNU: return decode_e8m0(value) target = result_type.mlir_type if value.type == target: return value if ir.IntegerType.isinstance(value.type): value = narrow_integer(value, result_type.width) return bitcast_to(value, target) def fragment_element(source, index: int) -> ir.Value: """Extract element `index` from a TensorSSA, or unwrap a scalar value.""" if isinstance(source, cute.TensorSSA): return vector.extract(source.ir_value(), [], [index]) if isinstance(source, ir.Value): return source if isinstance(source, (int, float)): return cutlass.Float32(source).ir_value() return source.ir_value() def zero_for_constraint(letter: str) -> ir.Value: """Create a zero in the register type named by an inline-asm constraint.""" target = CONSTRAINT_TYPES[letter]() return arith.constant(target, 0.0 if letter in ("f", "d") else 0) def packed_operands( sources, input_letters: list[str], scalar_sources: tuple[bool, ...], base: int, count: int, pack: int, ) -> list[ir.Value]: """Gather one source-major inline-asm pack, padding a partial tail with zero.""" operands = [] for source_index, source in enumerate(sources): for lane in range(pack): letter = input_letters[source_index * pack + lane] index = base + lane operands.append( bitcast_to( fragment_element( source, 0 if scalar_sources[source_index] else index ) if index < count else zero_for_constraint(letter), CONSTRAINT_TYPES[letter](), ) ) return operands @dsl_user_op def inline_asm_elementwise_intrinsic( *sources, asm: str, constraints: str, result_type, is_pure: bool = True, pack: int = 1, scalar_sources: tuple[bool, ...] = (), loc=None, ip=None, ): """Apply an inline PTX block elementwise across a TensorSSA fragment. Args: sources: Broadcastable TensorSSA fragments or repeated scalar values. asm: PTX text using `$N` operand syntax, output first. constraints: LLVM constraint list, e.g. `"=h,r"`. result_type: Logical cutlass numeric type of the produced fragment. E8M0 integer results are decoded to Float32 for fused consumers. is_pure: Whether the block may be reordered and CSEd. pack: Elements consumed per asm invocation. Missing tail elements are zero-padded and their corresponding outputs are discarded. scalar_sources: Sources whose physical TensorSSA value represents one logical scalar and must be repeated for every lane. """ if pack < 1: raise ValueError(f"CuteDSL inline asm requires pack >= 1, got pack={pack}") output_letters, input_letters = split_constraints(constraints) if len(output_letters) != pack: raise ValueError( f"inline asm pack={pack} requires {pack} output constraints, " f"got {len(output_letters)} in {constraints!r}" ) expected_inputs = len(sources) * pack if len(input_letters) != expected_inputs: raise ValueError( f"inline asm pack={pack} with {len(sources)} sources requires " f"{expected_inputs} input constraints, got {len(input_letters)} " f"in {constraints!r}" ) if not scalar_sources: scalar_sources = (False,) * len(sources) elif len(scalar_sources) != len(sources): raise ValueError( f"inline asm received {len(scalar_sources)} scalar-source flags " f"for {len(sources)} sources" ) output_types = [CONSTRAINT_TYPES[letter]() for letter in output_letters] asm_result_type = output_types[0] if len(output_types) > 1: fields = ", ".join(str(output_type) for output_type in output_types) asm_result_type = ir.Type.parse(f"!llvm.struct<({fields})>") compute_type = ( cutlass.Float32 if result_type == cutlass.Float8E8M0FNU else result_type ) fragments = [ source for source, is_scalar in zip(sources, scalar_sources) if isinstance(source, cute.TensorSSA) and not is_scalar ] if fragments: shape = fragments[0].shape if any(source.shape != shape for source in fragments[1:]): shape = _infer_broadcast_shape(*(source.shape for source in fragments)) sources = tuple( source.broadcast_to(shape) if isinstance(source, cute.TensorSSA) and not is_scalar else source for source, is_scalar in zip(sources, scalar_sources) ) first_fragment = next( source for source, is_scalar in zip(sources, scalar_sources) if isinstance(source, cute.TensorSSA) and not is_scalar ) count = math.prod(ir.VectorType(first_fragment.ir_value().type).shape) else: count = 1 shape = None converted = [] for base in range(0, count, pack): produced = llvm.inline_asm( asm_result_type, packed_operands(sources, input_letters, scalar_sources, base, count, pack), asm, constraints, has_side_effects=not is_pure, is_align_stack=False, ) outputs = [produced] if len(output_types) > 1: outputs = [ llvm.extractvalue(output_type, produced, [index]) for index, output_type in enumerate(output_types) ] valid_outputs = min(pack, count - base) converted.extend( convert_output(output, result_type) for output in outputs[:valid_outputs] ) if shape is None: return compute_type(converted[0]) vector_type = ir.VectorType.get([count], compute_type.mlir_type) return cute.TensorSSA( vector.from_elements(vector_type, converted), shape, compute_type )