Add arch specific target support

numba_cuda/numba/cuda/codegen.py

@@ -209,7 +209,9 @@ def _ensure_cc(self, cc):
             return cc
 
         device = devices.get_context().device
-        return device.compute_capability
+        cc = device.compute_capability
+        cc = (cc[0], cc[1], "a" if cc >= (9, 0) else "")
+        return cc
 
     def get_asm_str(self, cc=None):
         return "\n".join(self.get_asm_strs(cc=cc))
@@ -347,7 +349,7 @@ def get_cufunc(self):
         cufunc = self._cufunc_cache.get(device.id, None)
         if cufunc:
             return cufunc
-        cubin = self.get_cubin(cc=device.compute_capability)
+        cubin = self.get_cubin()
         module = ctx.create_module_image(
             cubin, self._setup_functions, self._teardown_functions
         )

numba_cuda/numba/cuda/cudadrv/driver.py

@@ -2303,7 +2303,11 @@ def __init__(
         lto=None,
         additional_flags=None,
     ):
-        arch = f"sm_{cc[0]}{cc[1]}"
+        if len(cc) == 3:
+            arch = f"sm_{cc[0]}{cc[1]}{cc[2]}"
+        else:
+            arch = f"sm_{cc[0]}{cc[1]}"
+
         self.max_registers = max_registers if max_registers else None
         self.lineinfo = lineinfo
         self.cc = cc

numba_cuda/numba/cuda/cudadrv/nvrtc.py

@@ -30,6 +30,44 @@ def _get_nvrtc_version():
     return (major, minor)
 
 
+def _verify_cc_tuple(cc):
+    version = _get_nvrtc_version()
+    ver_str = lambda version: ".".join(str(v) for v in version)
+
+    if len(cc) == 3:
+        cc, arch = (cc[0], cc[1]), cc[2]
+    else:
+        arch = ""
+
+    if arch not in ("", "a", "f"):
+        raise ValueError(
+            f"Invalid architecture suffix '{arch}' in compute capability "
+            f"{ver_str(cc)}{arch}. Expected '', 'a', or 'f'."
+        )
+
+    supported_ccs = get_supported_ccs()
+    try:
+        found = max(filter(lambda v: v <= cc, [v for v in supported_ccs]))
+    except ValueError:
+        raise RuntimeError(
+            f"Device compute capability {ver_str(cc)} is less than the "
+            f"minimum supported by NVRTC {ver_str(version)}. Supported "
+            "compute capabilities are "
+            f"{', '.join([ver_str(v) for v in supported_ccs])}."
+        )
+
+    if found != cc:
+        found = (found[0], found[1], arch)
+        warnings.warn(
+            f"Device compute capability {ver_str(cc)} is not supported by "
+            f"NVRTC {ver_str(version)}. Using {ver_str(found)} instead."
+        )
+    else:
+        found = (cc[0], cc[1], arch)
+
+    return found
+
+
 def compile(src, name, cc, ltoir=False, lineinfo=False, debug=False):
     """
     Compile a CUDA C/C++ source to PTX or LTOIR for a given compute capability.
@@ -38,7 +76,8 @@ def compile(src, name, cc, ltoir=False, lineinfo=False, debug=False):
     :type src: str
     :param name: The filename of the source (for information only)
     :type name: str
-    :param cc: A tuple ``(major, minor)`` of the compute capability
+    :param cc: A tuple ``(major, minor)`` or ``(major, minor, arch)`` of the
+        compute capability
     :type cc: tuple
     :param ltoir: Compile into LTOIR if True, otherwise into PTX
     :type ltoir: bool
@@ -49,34 +88,18 @@ def compile(src, name, cc, ltoir=False, lineinfo=False, debug=False):
     :return: The compiled PTX or LTOIR and compilation log
     :rtype: tuple
     """
+    found = _verify_cc_tuple(cc)
     version = _get_nvrtc_version()
 
-    ver_str = lambda version: ".".join(str(v) for v in version)
-    supported_ccs = get_supported_ccs()
-    try:
-        found = max(filter(lambda v: v <= cc, [v for v in supported_ccs]))
-    except ValueError:
-        raise RuntimeError(
-            f"Device compute capability {ver_str(cc)} is less than the "
-            f"minimum supported by NVRTC {ver_str(version)}. Supported "
-            "compute capabilities are "
-            f"{', '.join([ver_str(v) for v in supported_ccs])}."
-        )
-
-    if found != cc:
-        warnings.warn(
-            f"Device compute capability {ver_str(cc)} is not supported by "
-            f"NVRTC {ver_str(version)}. Using {ver_str(found)} instead."
-        )
-
     # Compilation options:
     # - Compile for the current device's compute capability.
     # - The CUDA include path is added.
     # - Relocatable Device Code (rdc) is needed to prevent device functions
     #   being optimized away.
-    major, minor = found
+    major, minor = found[0], found[1]
+    cc_arch = found[2] if len(found) == 3 else ""
 
-    arch = f"sm_{major}{minor}"
+    arch = f"sm_{major}{minor}{cc_arch}"
 
     cuda_include_dir = get_cuda_paths()["include_dir"].info
     cuda_includes = [f"{cuda_include_dir}"]
@@ -156,7 +179,7 @@ def write(self, msg):
     return result, log
 
 
-def find_closest_arch(mycc):
+def find_closest_arch(cc):
     """
     Given a compute capability, return the closest compute capability supported
     by the CUDA toolkit.
@@ -166,17 +189,17 @@ def find_closest_arch(mycc):
     """
     supported_ccs = get_supported_ccs()
 
-    for i, cc in enumerate(supported_ccs):
-        if cc == mycc:
+    for i, supported_cc in enumerate(supported_ccs):
+        if supported_cc == cc:
             # Matches
-            return cc
-        elif cc > mycc:
+            return supported_cc
+        elif supported_cc > cc:
             # Exceeded
             if i == 0:
                 # CC lower than supported
                 msg = (
                     "GPU compute capability %d.%d is not supported"
-                    "(requires >=%d.%d)" % (mycc + cc)
+                    "(requires >=%d.%d)" % (cc + supported_cc)
                 )
                 raise CCSupportError(msg)
             else:
@@ -187,13 +210,13 @@ def find_closest_arch(mycc):
     return supported_ccs[-1]  # Choose the highest
 
 
-def get_arch_option(major, minor):
+def get_arch_option(major, minor, arch=""):
     """Matches with the closest architecture option"""
     if config.FORCE_CUDA_CC:
-        arch = config.FORCE_CUDA_CC
+        major, minor = config.FORCE_CUDA_CC
     else:
-        arch = find_closest_arch((major, minor))
-    return "compute_%d%d" % arch
+        major, minor = find_closest_arch((major, minor))
+    return f"compute_{major}{minor}{arch}"
 
 
 def get_lowest_supported_cc():

numba_cuda/numba/cuda/testing.py

@@ -300,6 +300,10 @@ def skip_unless_cc_75(fn):
     return unittest.skipUnless(cc_X_or_above(7, 5), "requires cc >= 7.5")(fn)
 
 
+def skip_unless_cc_90(fn):
+    return unittest.skipUnless(cc_X_or_above(9, 0), "requires cc >= 9.0")(fn)
+
+
 def xfail_unless_cudasim(fn):
     if config.ENABLE_CUDASIM:
         return fn

numba_cuda/numba/cuda/tests/cudadrv/test_cuda_driver.py

@@ -2,6 +2,8 @@
 # SPDX-License-Identifier: BSD-2-Clause
 
 from ctypes import c_int, sizeof
+import cffi
+import numpy as np
 
 from numba.cuda.cudadrv.driver import host_to_device, device_to_host, driver
 from numba.cuda._compat import (
@@ -13,10 +15,12 @@
 
 from numba import cuda
 from numba.cuda.cudadrv import devices
-from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import unittest, CUDATestCase, skip_unless_cc_90
 from numba.cuda.testing import skip_on_cudasim
+from numba.core import types
 import contextlib
 
+
 ptx1 = """
     .version 1.4
     .target sm_10, map_f64_to_f32
@@ -390,5 +394,46 @@ def test_device_get_uuid(self):
         self.assertRegex(dev.uuid, uuid_format)
 
 
+@skip_unless_cc_90
+@skip_on_cudasim("CUDA asm unsupported in the simulator")
+class TestAcceleratedArchitecture(CUDATestCase):
+    def test_device_arch_specific(self):
+        set_desc = cuda.CUSource("""
+        #include <cuda_fp16.h>
+
+        extern "C" __device__
+        int set_descriptor(int *out, int* smem) {
+            unsigned usmem = __cvta_generic_to_shared(smem);
+            asm volatile("tensormap.replace.tile.rank.shared::cta.b1024.b32 [%0], 2;" :: "r"(usmem));
+            return 0;
+        }
+        """)
+
+        set_descriptor = cuda.declare_device(
+            "set_descriptor",
+            types.int32(types.CPointer(types.int32)),
+            link=[set_desc],
+        )
+
+        ffi = cffi.FFI()
+
+        @cuda.jit
+        def kernel(a):
+            sm = cuda.shared.array(1, dtype=np.int32)
+            data_ptr = ffi.from_buffer(sm)
+            set_descriptor(data_ptr)
+
+            # just to prevent optimization:
+            sm[0] = 2
+            cuda.syncthreads()
+            a[0] = sm[0]
+
+        a = np.ones(1, dtype=np.int32)
+
+        kernel[1, 1](a)
+
+        assert a[0] == 2
+
+
 if __name__ == "__main__":
     unittest.main()