fix: warp vote operations must use a constant int for the mode parameter (#606)

Fixes #592. Followed similar pattern as
#231.

Author: kaeun97

numba_cuda/numba/cuda/cudadecl.py

@@ -103,16 +103,6 @@ class Cuda_syncwarp(ConcreteTemplate):
     cases = [signature(types.none), signature(types.none, types.i4)]
 
 
-@register
-class Cuda_vote_sync_intrinsic(ConcreteTemplate):
-    key = cuda.vote_sync_intrinsic
-    cases = [
-        signature(
-            types.Tuple((types.i4, types.b1)), types.i4, types.i4, types.b1
-        )
-    ]
-
-
 @register
 class Cuda_match_any_sync(ConcreteTemplate):
     key = cuda.match_any_sync
@@ -522,9 +512,6 @@ def resolve_threadfence_system(self, mod):
     def resolve_syncwarp(self, mod):
         return types.Function(Cuda_syncwarp)
 
-    def resolve_vote_sync_intrinsic(self, mod):
-        return types.Function(Cuda_vote_sync_intrinsic)
-
     def resolve_match_any_sync(self, mod):
         return types.Function(Cuda_match_any_sync)
 

numba_cuda/numba/cuda/cudaimpl.py

@@ -280,18 +280,6 @@ def ptx_syncwarp_mask(context, builder, sig, args):
     return context.get_dummy_value()
 
 
-@lower(stubs.vote_sync_intrinsic, types.i4, types.i4, types.boolean)
-def ptx_vote_sync(context, builder, sig, args):
-    fname = "llvm.nvvm.vote.sync"
-    lmod = builder.module
-    fnty = ir.FunctionType(
-        ir.LiteralStructType((ir.IntType(32), ir.IntType(1))),
-        (ir.IntType(32), ir.IntType(32), ir.IntType(1)),
-    )
-    func = cgutils.get_or_insert_function(lmod, fnty, fname)
-    return builder.call(func, args)
-
-
 @lower(stubs.match_any_sync, types.i4, types.i4)
 @lower(stubs.match_any_sync, types.i4, types.i8)
 @lower(stubs.match_any_sync, types.i4, types.f4)

numba_cuda/numba/cuda/device_init.py

@@ -27,7 +27,6 @@
     local,
     const,
     atomic,
-    vote_sync_intrinsic,
     match_any_sync,
     match_all_sync,
     threadfence_block,
@@ -56,6 +55,10 @@
     shfl_up_sync,
     shfl_down_sync,
     shfl_xor_sync,
+    all_sync,
+    any_sync,
+    eq_sync,
+    ballot_sync,
 )
 from .cudadrv.error import CudaSupportError
 from numba.cuda.cudadrv.driver import (
@@ -79,12 +82,6 @@
 from .api import _auto_device
 from .args import In, Out, InOut
 
-from .intrinsic_wrapper import (
-    all_sync,
-    any_sync,
-    eq_sync,
-    ballot_sync,
-)
 
 from .kernels import reduction
 from numba.cuda.cudadrv.linkable_code import (

numba_cuda/numba/cuda/intrinsic_wrapper.py

@@ -6,7 +6,11 @@
 from numba import cuda
 from numba.cuda import types
 from numba.cuda import cgutils
-from numba.cuda.core.errors import RequireLiteralValue, TypingError
+from numba.cuda.core.errors import (
+    RequireLiteralValue,
+    TypingError,
+    NumbaTypeError,
+)
 from numba.cuda.typing import signature
 from numba.cuda.extending import overload_attribute, overload_method
 from numba.cuda import nvvmutils
@@ -380,3 +384,148 @@ def codegen(context, builder, sig, args):
     sig = signature(a_type, membermask_type, a_type, b_type)
 
     return sig, codegen
+
+
+# -------------------------------------------------------------------------------
+# Warp vote functions
+#
+# References:
+#
+# - https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-vote-functions
+# - https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html?highlight=data%2520movement#vote
+#
+# Notes:
+#
+# - The NVVM IR specification requires some of the mode parameter to be
+#   constants. It's therefore essential that we pass in mode values to the
+#   vote_sync_intrinsic.
+
+
+@intrinsic
+def all_sync(typingctx, mask_type, predicate_type):
+    """
+    If for all threads in the masked warp the predicate is true, then
+    a non-zero value is returned, otherwise 0 is returned.
+    """
+    mode_value = 0
+    sig, codegen_inner = vote_sync_intrinsic(
+        typingctx, mask_type, mode_value, predicate_type
+    )
+
+    def codegen(context, builder, sig_outer, args):
+        # Call vote_sync_intrinsic and extract the boolean result (index 1)
+        result_tuple = codegen_inner(context, builder, sig, args)
+        return builder.extract_value(result_tuple, 1)
+
+    sig_outer = signature(types.b1, mask_type, predicate_type)
+    return sig_outer, codegen
+
+
+@intrinsic
+def any_sync(typingctx, mask_type, predicate_type):
+    """
+    If for any thread in the masked warp the predicate is true, then
+    a non-zero value is returned, otherwise 0 is returned.
+    """
+    mode_value = 1
+    sig, codegen_inner = vote_sync_intrinsic(
+        typingctx, mask_type, mode_value, predicate_type
+    )
+
+    def codegen(context, builder, sig_outer, args):
+        result_tuple = codegen_inner(context, builder, sig, args)
+        return builder.extract_value(result_tuple, 1)
+
+    sig_outer = signature(types.b1, mask_type, predicate_type)
+    return sig_outer, codegen
+
+
+@intrinsic
+def eq_sync(typingctx, mask_type, predicate_type):
+    """
+    If for all threads in the masked warp the boolean predicate is the same,
+    then a non-zero value is returned, otherwise 0 is returned.
+    """
+    mode_value = 2
+    sig, codegen_inner = vote_sync_intrinsic(
+        typingctx, mask_type, mode_value, predicate_type
+    )
+
+    def codegen(context, builder, sig_outer, args):
+        result_tuple = codegen_inner(context, builder, sig, args)
+        return builder.extract_value(result_tuple, 1)
+
+    sig_outer = signature(types.b1, mask_type, predicate_type)
+    return sig_outer, codegen
+
+
+@intrinsic
+def ballot_sync(typingctx, mask_type, predicate_type):
+    """
+    Returns a mask of all threads in the warp whose predicate is true,
+    and are within the given mask.
+    """
+    mode_value = 3
+    sig, codegen_inner = vote_sync_intrinsic(
+        typingctx, mask_type, mode_value, predicate_type
+    )
+
+    def codegen(context, builder, sig_outer, args):
+        result_tuple = codegen_inner(context, builder, sig, args)
+        return builder.extract_value(
+            result_tuple, 0
+        )  # Extract ballot result (index 0)
+
+    sig_outer = signature(types.i4, mask_type, predicate_type)
+    return sig_outer, codegen
+
+
+def vote_sync_intrinsic(typingctx, mask_type, mode_value, predicate_type):
+    # Validate mode value
+    if mode_value not in (0, 1, 2, 3):
+        raise ValueError("Mode must be 0 (all), 1 (any), 2 (eq), or 3 (ballot)")
+
+    if types.unliteral(mask_type) not in types.integer_domain:
+        raise NumbaTypeError(f"Mask type must be an integer. Got {mask_type}")
+    predicate_types = types.integer_domain | {types.boolean}
+
+    if types.unliteral(predicate_type) not in predicate_types:
+        raise NumbaTypeError(
+            f"Predicate must be an integer or boolean. Got {predicate_type}"
+        )
+
+    def codegen(context, builder, sig, args):
+        mask, predicate = args
+
+        # Types
+        i1 = ir.IntType(1)
+        i32 = ir.IntType(32)
+
+        # NVVM intrinsic definition
+        arg_types = (i32, i32, i1)
+        vote_return_type = ir.LiteralStructType((i32, i1))
+        fnty = ir.FunctionType(vote_return_type, arg_types)
+
+        fname = "llvm.nvvm.vote.sync"
+        lmod = builder.module
+        vote_sync = cgutils.get_or_insert_function(lmod, fnty, fname)
+
+        # Intrinsic arguments
+        mode = ir.Constant(i32, mode_value)
+        mask_i32 = builder.trunc(mask, i32)
+
+        # Convert predicate to i1
+        if predicate.type != ir.IntType(1):
+            predicate_bool = builder.icmp_signed(
+                "!=", predicate, ir.Constant(predicate.type, 0)
+            )
+        else:
+            predicate_bool = predicate
+
+        return builder.call(vote_sync, [mask_i32, mode, predicate_bool])
+
+    sig = signature(
+        types.Tuple((types.i4, types.b1)), mask_type, predicate_type
+    )
+
+    return sig, codegen

numba_cuda/numba/cuda/intrinsics.py

@@ -200,17 +200,6 @@ class syncwarp(Stub):
     _description_ = "<warp_sync()>"
 
 
-class vote_sync_intrinsic(Stub):
-    """
-    vote_sync_intrinsic(mask, mode, predictate)
-
-    Nvvm intrinsic for performing a reduce and broadcast across a warp
-    docs.nvidia.com/cuda/nvvm-ir-spec/index.html#nvvm-intrin-warp-level-vote
-    """
-
-    _description_ = "<vote_sync()>"
-
-
 class match_any_sync(Stub):
     """
     match_any_sync(mask, value)