|
| 1 | +# SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 2 | +# |
| 3 | +# SPDX-License-Identifier: Apache-2.0 |
| 4 | + |
| 5 | +from dataclasses import dataclass |
| 6 | +from enum import Enum |
| 7 | +from typing import Any, Callable, Literal, TypeVar |
| 8 | + |
| 9 | +import cuda.lang._datatype as datatype |
| 10 | +from cuda.lang._execution import stub |
| 11 | +import cuda.lang._mlir as mlir |
| 12 | +from cuda.lang._stub._nvvm_support import ( |
| 13 | + _IntrinsicDTypeAnnotation, |
| 14 | + _IntrinsicPredicateAnnotation, |
| 15 | +) |
| 16 | +from cuda.lang._ir.type import TileTy |
| 17 | +from cuda.tile import TileTypeError, TileValueError |
| 18 | +from cuda.tile._ir.op_impl import ( |
| 19 | + require_constant_bool, |
| 20 | + require_constant_enum, |
| 21 | + require_constant_int, |
| 22 | +) |
| 23 | +from cuda.tile._ir.ir import Var, add_operation_variadic |
| 24 | +from cuda.tile._ir.ops import ( |
| 25 | + implicit_cast, |
| 26 | + build_tuple, |
| 27 | +) |
| 28 | + |
| 29 | + |
| 30 | +FuncTy = TypeVar("FuncTy", bound=Callable[..., Any]) |
| 31 | + |
| 32 | + |
| 33 | +@dataclass(frozen=True) |
| 34 | +class ArgSpec: |
| 35 | + type: object |
| 36 | + kind: Literal["operand", "attribute"] = "operand" |
| 37 | + optional: bool = False |
| 38 | + variadic: bool = False |
| 39 | + unit: bool = False |
| 40 | + name: str = "" |
| 41 | + |
| 42 | + |
| 43 | +@dataclass(frozen=True) |
| 44 | +class ResultSpec: |
| 45 | + name: str |
| 46 | + type: object |
| 47 | + optional: bool = False |
| 48 | + variadic: bool = False |
| 49 | + |
| 50 | + |
| 51 | +def is_none_constant(value: Var) -> bool: |
| 52 | + return value.is_constant() and value.get_constant() is None |
| 53 | + |
| 54 | + |
| 55 | +def is_enum_type(ty) -> bool: |
| 56 | + return isinstance(ty, type) and issubclass(ty, Enum) |
| 57 | + |
| 58 | + |
| 59 | +def cast_operand(spec: ArgSpec, arg: Var) -> Var: |
| 60 | + target_type = spec.type |
| 61 | + src_type = arg.get_type() |
| 62 | + ctx = f"Attempting to cast argument to {target_type=}" |
| 63 | + match target_type: |
| 64 | + case _IntrinsicPredicateAnnotation(): |
| 65 | + target_type.predicate(arg) |
| 66 | + return arg |
| 67 | + case _IntrinsicDTypeAnnotation(): |
| 68 | + return implicit_cast(arg, target_type.dtype, ctx) |
| 69 | + case tuple(): |
| 70 | + for target in target_type: |
| 71 | + try: |
| 72 | + return implicit_cast(arg, target.dtype, ctx) |
| 73 | + except (TileTypeError, TileValueError): |
| 74 | + pass |
| 75 | + options = ", ".join([str(t) for t in target_type]) |
| 76 | + raise TileTypeError( |
| 77 | + f"Could not cast arg of type {src_type} to any of {options}" |
| 78 | + ) |
| 79 | + case _: |
| 80 | + raise TileTypeError("Expected a predicate, a dtype, or a tuple of dtypes") |
| 81 | + |
| 82 | + |
| 83 | +def make_mlir_attribute(spec: ArgSpec, arg: Var) -> tuple[str, mlir.Attribute] | None: |
| 84 | + if spec.optional and is_none_constant(arg): |
| 85 | + return None |
| 86 | + |
| 87 | + if is_enum_type(spec.type): |
| 88 | + attr_cls = getattr(mlir.nvvm, spec.type.__name__ + "Attr") |
| 89 | + arg = require_constant_enum(arg, spec.type) |
| 90 | + return spec.name, attr_cls(value=arg) |
| 91 | + |
| 92 | + if spec.unit: |
| 93 | + arg = require_constant_bool(arg) |
| 94 | + return (spec.name, mlir.UnitAttr()) if arg else None |
| 95 | + |
| 96 | + dtype = spec.type.dtype |
| 97 | + if dtype is datatype.bool_: |
| 98 | + arg = require_constant_bool(arg) |
| 99 | + return spec.name, mlir.BoolAttr(value=arg) |
| 100 | + |
| 101 | + if datatype.is_integral(dtype): |
| 102 | + arg = require_constant_int(arg) |
| 103 | + ty = mlir.IntegerType.signless(dtype.bitwidth) |
| 104 | + attr = mlir.IntegerAttr.make(ty, int(arg)) |
| 105 | + return spec.name, attr |
| 106 | + |
| 107 | + raise TileTypeError(f"Cannot convert argument into attribute: {spec}") |
| 108 | + |
| 109 | + |
| 110 | +def get_raw_mlir_parts( |
| 111 | + arg_specs, has_operand_segment_sizes, args: tuple[Var, ...] |
| 112 | +) -> tuple[tuple[Var, ...], tuple[tuple[str, mlir.Attribute], ...]]: |
| 113 | + operands = [] |
| 114 | + attributes = [] |
| 115 | + operand_segment_sizes = [] |
| 116 | + for arg, spec in zip(args, arg_specs, strict=True): |
| 117 | + if spec.kind == "attribute": |
| 118 | + attr = make_mlir_attribute(spec, arg) |
| 119 | + if attr is not None: |
| 120 | + attributes.append(attr) |
| 121 | + |
| 122 | + elif spec.optional and is_none_constant(arg): |
| 123 | + operand_segment_sizes.append(0) |
| 124 | + |
| 125 | + elif spec.variadic: |
| 126 | + assert isinstance(arg, tuple) |
| 127 | + operands.extend(cast_operand(spec, item) for item in arg) |
| 128 | + operand_segment_sizes.append(len(arg)) |
| 129 | + else: |
| 130 | + operands.append(cast_operand(spec, arg)) |
| 131 | + operand_segment_sizes.append(1) |
| 132 | + |
| 133 | + if has_operand_segment_sizes: |
| 134 | + attributes.append( |
| 135 | + ("operandSegmentSizes", mlir.DenseI32ArrayAttr(operand_segment_sizes)) |
| 136 | + ) |
| 137 | + |
| 138 | + return tuple(operands), tuple(attributes) |
| 139 | + |
| 140 | + |
| 141 | +def _raw_nvvm_mlir_operation_impl(stub_func, *args: Var): |
| 142 | + from cuda.lang._ir.ops import RawMLIROperation |
| 143 | + |
| 144 | + result_types = tuple(TileTy(ty.type.dtype) for ty in stub_func._results) |
| 145 | + operands, attrs = get_raw_mlir_parts( |
| 146 | + stub_func._args, stub_func._attr_sized_operand_segments, args |
| 147 | + ) |
| 148 | + results = add_operation_variadic( |
| 149 | + RawMLIROperation, |
| 150 | + result_types, |
| 151 | + op_name=stub_func._op_name, |
| 152 | + operands_=operands, |
| 153 | + mlir_attributes=attrs, |
| 154 | + ) |
| 155 | + match len(stub_func._results): |
| 156 | + case 0: |
| 157 | + return None |
| 158 | + case 1: |
| 159 | + return results[0] |
| 160 | + case _: |
| 161 | + return build_tuple(results) |
| 162 | + |
| 163 | + |
| 164 | +_raw_nvvm_mlir_operation_impl._is_coroutine = False |
| 165 | + |
| 166 | + |
| 167 | +def nvvm_mlir_interface_stub( |
| 168 | + *, |
| 169 | + op_name: str, |
| 170 | + attr_sized_operand_segments: bool = False, |
| 171 | + results: tuple[ResultSpec, ...] = (), |
| 172 | + args: tuple[ArgSpec, ...] = (), |
| 173 | +) -> Callable[[FuncTy], FuncTy]: |
| 174 | + def decorate(func: FuncTy) -> FuncTy: |
| 175 | + func = stub(func) |
| 176 | + func._cutile_custom_implementation_handler = _raw_nvvm_mlir_operation_impl |
| 177 | + func._op_name = op_name |
| 178 | + func._attr_sized_operand_segments = attr_sized_operand_segments |
| 179 | + func._results = results |
| 180 | + func._args = args |
| 181 | + return func |
| 182 | + |
| 183 | + return decorate |
| 184 | + |
| 185 | + |
| 186 | +__all__ = ("ArgSpec", "ResultSpec", "nvvm_mlir_interface_stub") |
0 commit comments