cuda.parallel: Check compiled code for LDL/STL instructions in tests (#4472)

* cuda.parallel: Check compiled code for LDL/STL instructions in tests

* No need for global & use monkeypatch

---------

Co-authored-by: Ashwin Srinath <[email protected]>

Author: shwina

python/cuda_parallel/cuda/parallel/experimental/_bindings.pyx

@@ -963,7 +963,8 @@ cdef class CommonData:
 
 cdef extern from "cccl/c/reduce.h":
     cdef struct cccl_device_reduce_build_result_t 'cccl_device_reduce_build_result_t':
-        pass
+        const char* cubin
+        size_t cubin_size
 
     cdef CUresult cccl_device_reduce_build(
         cccl_device_reduce_build_result_t*,
@@ -1071,6 +1072,8 @@ cdef class DeviceReduceBuildResult:
             )
         return storage_sz
 
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]
 
 # ------------
 #   DeviceScan
@@ -1081,7 +1084,8 @@ cdef extern from "cccl/c/scan.h":
     ctypedef bint _Bool
 
     cdef struct cccl_device_scan_build_result_t 'cccl_device_scan_build_result_t':
-        pass
+        const char* cubin
+        size_t cubin_size
 
     cdef CUresult cccl_device_scan_build(
         cccl_device_scan_build_result_t*,
@@ -1236,6 +1240,8 @@ cdef class DeviceScanBuildResult:
             )
         return storage_sz
 
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]
 
 # -----------------------
 #   DeviceSegmentedReduce
@@ -1244,7 +1250,8 @@ cdef class DeviceScanBuildResult:
 
 cdef extern from "cccl/c/segmented_reduce.h":
     cdef struct cccl_device_segmented_reduce_build_result_t 'cccl_device_segmented_reduce_build_result_t':
-        pass
+        const char* cubin
+        size_t cubin_size
 
     cdef CUresult cccl_device_segmented_reduce_build(
         cccl_device_segmented_reduce_build_result_t*,
@@ -1364,14 +1371,17 @@ cdef class DeviceSegmentedReduceBuildResult:
             )
         return storage_sz
 
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]
 # -----------------
 #   DeviceMergeSort
 # -----------------
 
 
 cdef extern from "cccl/c/merge_sort.h":
     cdef struct cccl_device_merge_sort_build_result_t 'cccl_device_merge_sort_build_result_t':
-        pass
+        const char* cubin
+        size_t cubin_size
 
     cdef CUresult cccl_device_merge_sort_build(
         cccl_device_merge_sort_build_result_t *bld_ptr,
@@ -1484,13 +1494,20 @@ cdef class DeviceMergeSortBuildResult:
             )
         return storage_sz
 
+
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]
+
+
 # -------------------
 #   DeviceUniqueByKey
 # -------------------
 
 cdef extern from "cccl/c/unique_by_key.h":
     cdef struct cccl_device_unique_by_key_build_result_t 'cccl_device_unique_by_key_build_result_t':
-        pass
+        const char* cubin
+        size_t cubin_size
+
 
     cdef CUresult cccl_device_unique_by_key_build(
         cccl_device_unique_by_key_build_result_t *build_ptr,
@@ -1611,12 +1628,16 @@ cdef class DeviceUniqueByKeyBuildResult:
         return storage_sz
 
 
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]
+
 # --------------------------------------------
 #   DeviceUnaryTransform/DeviceBinaryTransform
 # --------------------------------------------
 cdef extern from "cccl/c/transform.h":
     cdef struct cccl_device_transform_build_result_t:
-        pass
+        const char* cubin
+        size_t cubin_size
 
     cdef CUresult cccl_device_unary_transform_build(
         cccl_device_transform_build_result_t *build_ptr,
@@ -1723,6 +1744,10 @@ cdef class DeviceUnaryTransform:
             raise RuntimeError("Failed to compute unary transform")
 
 
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]
+
+
 cdef class DeviceBinaryTransform:
     cdef cccl_device_transform_build_result_t build_data
 
@@ -1791,3 +1816,6 @@ cdef class DeviceBinaryTransform:
             )
         if (status != 0):
             raise RuntimeError("Failed to compute binary transform")
+
+    def _get_cubin(self):
+        return self.build_data.cubin[:self.build_data.cubin_size]

python/cuda_parallel/cuda/parallel/experimental/_cccl_interop.py

@@ -5,6 +5,9 @@
 from __future__ import annotations
 
 import functools
+import os
+import subprocess
+import tempfile
 from typing import Callable, List
 
 import numba
@@ -220,6 +223,29 @@ def get_paths() -> List[str]:
     return paths
 
 
+def _check_compile_result(cubin: bytes):
+    # check compiled code for LDL/STL instructions
+    temp_cubin_file = tempfile.NamedTemporaryFile(delete=False)
+    try:
+        temp_cubin_file.write(cubin)
+        out = subprocess.run(
+            ["nvdisasm", "-gi", temp_cubin_file.name], capture_output=True
+        )
+        if out.returncode != 0:
+            raise RuntimeError("nvdisasm failed")
+        sass = out.stdout.decode("utf-8")
+    finally:
+        os.unlink(temp_cubin_file.name)
+
+    assert "LDL" not in sass, "LDL instruction found in SASS"
+    assert "STL" not in sass, "STL instruction found in SASS"
+
+
+# this global variable controls whether the compile result is checked
+# for LDL/STL instructions. Should be set to `True` for testing only.
+_check_sass: bool = False
+
+
 def call_build(build_impl_fn: Callable, *args, **kwargs):
     """Calls given build_impl_fn callable while providing compute capability and paths
 
@@ -230,9 +256,12 @@ def call_build(build_impl_fn: Callable, *args, **kwargs):
     common_data = CommonData(
         cc_major, cc_minor, cub_path, thrust_path, libcudacxx_path, cuda_include_path
     )
-    error = build_impl_fn(
+    result = build_impl_fn(
         *args,
         common_data,
         **kwargs,
     )
-    return error
+    if _check_sass:
+        cubin = result._get_cubin()
+        _check_compile_result(cubin)
+    return result

python/cuda_parallel/tests/conftest.py

@@ -67,3 +67,14 @@ def ptr(self):
 @pytest.fixture(scope="function")
 def cuda_stream() -> Stream:
     return Stream(cp.cuda.Stream())
+
+
+@pytest.fixture(scope="function", autouse=True)
+def verify_sass(monkeypatch):
+    import cuda.parallel.experimental._cccl_interop
+
+    monkeypatch.setattr(
+        cuda.parallel.experimental._cccl_interop,
+        "_check_sass",
+        False,  # todo: change to True
+    )