[lang] Inline PTX support for LLVMIR backend

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

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

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import itertools
 from dataclasses import dataclass
 from typing import Optional, Any, TYPE_CHECKING
 
@@ -22,6 +23,7 @@
 from cuda.tile._memory_model import MemoryScope
 from cuda.tile._ir.ir import MemoryEffect, add_operation_variadic
 from cuda.tile._ir.type import TensorLikeTy
+from cuda.lang import _datatype as datatype
 from cuda.lang._enums import VectorReduction
 from .ir import Operation, Var, attribute, operand
 from .type import Type, VectorTy, ScalarTy, PointerTy
@@ -83,8 +85,7 @@ def generate_llvm(self, ctx):
         elif len(self.result_vars) == 1:
             return (result,)
         else:
-            return tuple(ctx.builder.extract_value(ctx.typeof(r), result, i)
-                         for i, r in enumerate(self.result_vars))
+            return ctx.unpack_struct(result)
 
 
 def call_intrinsic(stub, *args: Var):
@@ -200,6 +201,71 @@ class InlinePTX(Operation, opcode="inline_ptx", memory_effect=MemoryEffect.STORE
     text: tuple[InlineAsmPiece, ...] = attribute()
     inputs: tuple[Var, ...] = operand()
 
+    @override
+    def generate_llvm(self, ctx):
+        num_outputs = len(self.result_vars)
+        llvm_types = []
+        constraints = []
+        for i, var in enumerate(itertools.chain(self.result_vars, self.inputs)):
+            ty = var.get_type()
+            assert ty.tensor_shape() == ()
+            dtype = ty.tensor_dtype()
+            code = _dtype_to_inline_ptx_constraint(dtype)
+            prefix = "=" if i < num_outputs else ""
+            constraints.append(prefix + code)
+            llvm_types.append(ctx.typeof(var))
+
+        pieces = []
+        for p in self.text:
+            if isinstance(p, str):
+                pieces.append(p)
+            else:
+                if isinstance(p, InlineAsmInput):
+                    linear_index = num_outputs + p.index
+                else:
+                    assert isinstance(p, InlineAsmOutput)
+                    linear_index = p.index
+                pieces.append(f"${linear_index}")
+
+        tt = ctx.builder.type_table
+        if num_outputs == 0:
+            ret_ty = tt.VOID
+        elif num_outputs == 1:
+            ret_ty = llvm_types[0]
+        else:
+            ret_ty = tt.struct_anonymous(llvm_types[:num_outputs])
+
+        func_ty = tt.function(ret_ty, llvm_types[num_outputs:])
+
+        asm = ctx.builder.constants.inline_asm(func_ty, "".join(pieces), ",".join(constraints),
+                                               side_effects=True)
+        r = ctx.builder.call(func_ty, asm, [ctx.value(x) for x in self.inputs])
+
+        if num_outputs == 0:
+            return ()
+        elif num_outputs == 1:
+            return (r,)
+        else:
+            return ctx.unpack_struct(r)
+
+
+def _dtype_to_inline_ptx_constraint(dtype: datatype.DType) -> str:
+    if dtype == datatype.float32:
+        return "f"
+    elif dtype == datatype.float64:
+        return "d"
+    elif dtype == datatype.bool_:
+        return "b"
+    elif dtype.bitwidth == 16:
+        return "h"
+    elif dtype.bitwidth == 32:
+        return "r"
+    elif dtype.bitwidth == 64:
+        return "l"
+    elif dtype.bitwidth == 128:
+        return "q"
+    raise NotImplementedError(f"Can't map dtype {dtype} to InlineAsm constraint")
+
 
 @dataclass(eq=False)
 class Fence(Operation, opcode="fence", memory_effect=MemoryEffect.STORE):

experimental/cuda-lang/src/cuda/lang/_llvm_bitcode/__init__.py

@@ -12,11 +12,15 @@
 )
 
 from ._builder import (
+    AnonStructType,
     BitcodeBuilder,
+    FloatType,
     Function,
     FunctionType,
+    IntegerType,
     Metadata,
     MetadataTable,
+    PointerType,
     Type,
     TypeTable,
     Value,
@@ -36,11 +40,15 @@
     "CmpPredicate",
     "FloatKind",
     "Linkage",
+    "AnonStructType",
     "BitcodeBuilder",
+    "FloatType",
     "Function",
     "FunctionType",
+    "IntegerType",
     "Metadata",
     "MetadataTable",
+    "PointerType",
     "Type",
     "TypeTable",
     "Value",

experimental/cuda-lang/src/cuda/lang/_llvm_bitcode/_builder.py

@@ -366,7 +366,8 @@ def is_declaration(self) -> bool:
 @dataclass
 class GlobalVariable:
     name: str
-    value: Value
+    value: Value  # Has a pointer type
+    value_type: Type  # Actual type of the global
     address_space: int
     is_constant: bool
     initializer: Value | None
@@ -428,10 +429,11 @@ def global_variable(self,
                         linkage: Linkage = Linkage.External,
                         alignment: int | None = None) -> Value:
         assert self._cur_function is None, "Global variables must be put at the global scope"
-        value = Value(type)
+        value = Value(self.type_table.pointer(address_space))
         self._global_variables.append(GlobalVariable(
             name=name,
             value=value,
+            value_type=type,
             address_space=address_space,
             is_constant=is_constant,
             initializer=initializer,
@@ -524,9 +526,16 @@ def get_element_ptr(self, element_ty: Type, ptr: Value, *indices: Value) -> Valu
             *indices
         )
 
-    def extract_value(self, result_ty: Type, src: Value, *indices: int) -> Value:
+    def extract_value(self, src: Value, *indices: int) -> Value:
+        ty = src.type
+        for i in indices:
+            if isinstance(ty, AnonStructType):
+                ty = ty.fields[i]
+            else:
+                raise NotImplementedError()
+
         return self._instruction(
-            result_ty,
+            ty,
             codes.FUNC_CODE_INST_EXTRACTVAL,
             "V" + "i" * len(indices),
             src, *indices
@@ -804,7 +813,7 @@ def _write_global_variable_record(global_var: GlobalVariable, writer: _BitcodeWr
     writer.unabbrev_record(
         codes.MODULE_CODE_GLOBALVAR,
         *string_table[global_var.name.encode()],  # STRTAB offset & size
-        global_var.value.type.type_id,
+        global_var.value_type.type_id,
         (global_var.address_space << 2) | 2 | global_var.is_constant,
         0 if global_var.initializer is None else global_var.initializer.id + 1,
         global_var.linkage._value_,

experimental/cuda-lang/src/cuda/lang/_passes/ir2llvm.py

@@ -140,6 +140,11 @@ def declare_function(self, name: str, func_ty: llvm.FunctionType) -> llvm.Value:
             self._function_declarations[name] = fptr
         return fptr
 
+    def unpack_struct(self, struct_value: llvm.Value) -> tuple[llvm.Value, ...]:
+        assert isinstance(struct_value.type, llvm.AnonStructType)
+        return tuple(self.builder.extract_value(struct_value, i)
+                     for i in range(len(struct_value.type.fields)))
+
 
 def generate_nvvm_bitcode_for_kernel(body: ir.Region,
                                      symbol: str,

experimental/cuda-lang/test/test_nvvm_bitcode.py

@@ -87,8 +87,8 @@ def _make_inline_ptx_bitcode() -> bytes:
             side_effects=False
         )
         r = builder.call(ptx_functy, ptx_func, [a, b])
-        radd = builder.extract_value(tt.F32, r, 0)
-        rsub = builder.extract_value(tt.F32, r, 1)
+        radd = builder.extract_value(r, 0)
+        rsub = builder.extract_value(r, 1)
         builder.store(cp0, radd, alignment=4)
         builder.store(cp1, rsub, alignment=4)
         builder.ret()