[lang] Lower FP32 even vectors to packed NVVM operations

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

experimental/cuda-lang/src/cuda/lang/_compile.py

@@ -43,6 +43,7 @@
     ConstantConstraint,
 )
 from cuda.lang._exception import CompilerExecutionError
+from cuda.lang._target import TargetInfo
 from ._execution import kernel
 from cuda.lang._ir.ops import cuda_lang_impl_registry
 from ._ir._host_program import HostProgram, get_host_programs_by_var
@@ -283,18 +284,25 @@ def _dump(phase: str, contents: object) -> None:
     if log_flags.log_flattened_ir:
         _dump("Flattened IR", flattened_ir)
 
-    mlir_module = ir2mlir(signature, flattened_ir, ctx, compiler_options)
-    mlir_text = str(mlir_module)
-
-    if log_flags.log_mlir:
-        _dump("MLIR", mlir_text)
-
     if gpu_name is None or arch is None:
         cc = compute_capability or get_compute_capability()
         suffix = "a" if cc >= (9, 0) else ""
         gpu_name = gpu_name or cc.gpu_name + suffix
         arch = arch or cc.arch + suffix
 
+    target_info = TargetInfo.from_arch(arch)
+    mlir_module = ir2mlir(
+        signature,
+        flattened_ir,
+        ctx,
+        compiler_options,
+        target_info,
+    )
+    mlir_text = str(mlir_module)
+
+    if log_flags.log_mlir:
+        _dump("MLIR", mlir_text)
+
     need_nvvm = log_flags.log_nvvm or keep_nvvm
     need_ptx = log_flags.log_ptx or keep_ptx
     ptx_compiler_options = compiler_options._ptx_compiler_options

experimental/cuda-lang/src/cuda/lang/_ir/op_defs.py

@@ -7,7 +7,7 @@
 from enum import Enum, auto
 
 import cuda.lang._mlir as mlir
-from cuda.lang._enums import MemoryOrder
+from cuda.lang._enums import MemoryOrder, RoundingMode, SaturationMode
 from cuda.tile._memory_model import MemoryScope
 from cuda.tile._ir.ir import MemoryEffect
 import cuda.lang._datatype as datatype
@@ -143,3 +143,15 @@ class ReinterpretPointerAsArray(Operation, opcode="reinterpret_ptr_as_array"):
 @dataclass
 class TensorMapAsOpaquePtr(Operation, opcode="tensor_map_as_opaque_ptr"):
     tensor_map: Var = operand()
+
+
+@dataclass(eq=False)
+class FmaOperation(Operation, opcode="fma"):
+    x: Var = operand()
+    y: Var = operand()
+    z: Var = operand()
+    rounding_mode: RoundingMode = attribute()
+    saturation_mode: SaturationMode = attribute()
+    flush_to_zero: bool = attribute()
+    relu: bool = attribute()
+    oob: bool = attribute()

experimental/cuda-lang/src/cuda/lang/_ir/op_impl/math_impl.py

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import enum
 import operator
 
 from cuda.tile._ir.ir import add_operation_variadic
@@ -22,6 +21,7 @@
     RawNVVMIntrinsic,
     RawMLIROperation,
     ForeignFunction,
+    FmaOperation,
 )
 from cuda.lang._ir.type_checking_helpers import (
     broadcast_to_same_shape,
@@ -58,15 +58,6 @@
 impl = _registry.impl
 
 
-# TODO(ajm): need to bump llvm bindings to get this enum
-class _FmaSaturationMode(enum.Enum):
-    NONE = 0
-    SAT = 1
-
-    def _print_mlir_unqualified(self, printer):
-        printer(("none", "sat")[self.value])
-
-
 def math_impl_registry() -> ImplRegistry:
     return _registry
 
@@ -111,12 +102,12 @@ def math_fma_impl(
     y = promote_and_broadcast_to(y, ty)
     z = promote_and_broadcast_to(z, ty)
 
-    rounding_modes = {
-        RoundingMode.RM: mlir.nvvm.FPRoundingMode.RM,
-        RoundingMode.RN: mlir.nvvm.FPRoundingMode.RN,
-        RoundingMode.RP: mlir.nvvm.FPRoundingMode.RP,
-        RoundingMode.RZ: mlir.nvvm.FPRoundingMode.RZ,
-    }
+    rounding_modes = (
+        RoundingMode.RM,
+        RoundingMode.RN,
+        RoundingMode.RP,
+        RoundingMode.RZ,
+    )
     if rounding_mode not in rounding_modes:
         valid = ", ".join(str(mode) for mode in rounding_modes)
         raise TypeCheckingError(
@@ -127,25 +118,17 @@ def math_fma_impl(
         raise TypeCheckingError(
             "fma does not implement SaturationMode.SATFINITE yet"
         )
-    saturation_modes = {
-        SaturationMode.NONE: _FmaSaturationMode.NONE,
-        SaturationMode.SAT: _FmaSaturationMode.SAT,
-    }
-
-    rns = mlir.nvvm.FPRoundingModeAttr(value=rounding_modes[rounding_mode])
-    sat = mlir.nvvm.SaturationModeAttr(value=saturation_modes[saturation_mode])
     return add_operation(
-        RawMLIROperation,
+        FmaOperation,
         ty,
-        op_name="nvvm.fma",
-        operands_=(x, y, z),
-        mlir_attributes=(
-            ("rnd", rns),
-            ("sat", sat),
-            ("ftz", mlir.BoolAttr(value=flush_to_zero)),
-            ("relu", mlir.BoolAttr(value=relu)),
-            ("oob", mlir.BoolAttr(value=oob)),
-        ),
+        x=x,
+        y=y,
+        z=z,
+        rounding_mode=rounding_mode,
+        saturation_mode=saturation_mode,
+        flush_to_zero=flush_to_zero,
+        relu=relu,
+        oob=oob,
     )
 
 

experimental/cuda-lang/src/cuda/lang/_ir/ops.py

@@ -101,6 +101,7 @@
     LoadPointer,
     ReinterpretPointerAsArray,
     BitCast,
+    FmaOperation,
 )
 from .op_impl.core_api_impl import core_api_impl_registry
 from .type_checking_helpers import (
@@ -942,6 +943,7 @@ def _call_foreign_function_impl(func: Var, return_type: Var, parameters: Var):
     "RawMLIROperation",
     "Fence",
     "ForeignFunction",
+    "FmaOperation",
     "VectorGetItem",
     "VectorReduce",
 )