Fix for atomic load, store, exchange failure on ocl (#1336)

This PR fixes the failure that was occurring `opencl:cpu` when atomic
load, store and exchange operations were being used. Spirv expects
arguments to be of integer types always. This PR introduces a `bitcast`
to integer types when the argument are of floating point types.
This PR also added additional function attributes to these atomic
function declarations, keeping it the same as the declaration generated
for DPC++.
It further updates the test cases to remove expected failure on
`opencl:cpu` executions.

Author: Diptorup Deb

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_atomic_ref_overloads.py

@@ -231,27 +231,52 @@ def _intrinsic_load(
 
     def _intrinsic_load_gen(context, builder, sig, args):
         atomic_ref_ty = sig.args[0]
-        fn = get_or_insert_atomic_load_fn(
-            context, builder.module, atomic_ref_ty
-        )
 
-        spirv_memory_semantics_mask = get_memory_semantics_mask(
-            atomic_ref_ty.memory_order
+        atomic_ref_ptr = builder.extract_value(
+            args[0],
+            context.data_model_manager.lookup(atomic_ref_ty).get_field_position(
+                "ref"
+            ),
         )
-        spirv_scope = get_scope(atomic_ref_ty.memory_scope)
+        if sig.args[0].dtype == types.float32:
+            atomic_ref_ptr = builder.bitcast(
+                atomic_ref_ptr,
+                llvmir.PointerType(
+                    llvmir.IntType(32), addrspace=sig.args[0].address_space
+                ),
+            )
+        elif sig.args[0].dtype == types.float64:
+            atomic_ref_ptr = builder.bitcast(
+                atomic_ref_ptr,
+                llvmir.PointerType(
+                    llvmir.IntType(64), addrspace=sig.args[0].address_space
+                ),
+            )
 
         fn_args = [
-            builder.extract_value(
-                args[0],
-                context.data_model_manager.lookup(
-                    atomic_ref_ty
-                ).get_field_position("ref"),
+            atomic_ref_ptr,
+            context.get_constant(
+                types.int32, get_scope(atomic_ref_ty.memory_scope)
+            ),
+            context.get_constant(
+                types.int32,
+                get_memory_semantics_mask(atomic_ref_ty.memory_order),
             ),
-            context.get_constant(types.int32, spirv_scope),
-            context.get_constant(types.int32, spirv_memory_semantics_mask),
         ]
 
-        return builder.call(fn, fn_args)
+        ret_val = builder.call(
+            get_or_insert_atomic_load_fn(
+                context, builder.module, atomic_ref_ty
+            ),
+            fn_args,
+        )
+
+        if sig.args[0].dtype == types.float32:
+            ret_val = builder.bitcast(ret_val, llvmir.FloatType())
+        elif sig.args[0].dtype == types.float64:
+            ret_val = builder.bitcast(ret_val, llvmir.DoubleType())
+
+        return ret_val
 
     return sig, _intrinsic_load_gen
 
@@ -264,28 +289,49 @@ def _intrinsic_store(
 
     def _intrinsic_store_gen(context, builder, sig, args):
         atomic_ref_ty = sig.args[0]
-        atomic_store_fn = get_or_insert_spv_atomic_store_fn(
-            context, builder.module, atomic_ref_ty
+
+        store_arg = args[1]
+        atomic_ref_ptr = builder.extract_value(
+            args[0],
+            context.data_model_manager.lookup(atomic_ref_ty).get_field_position(
+                "ref"
+            ),
         )
+        if sig.args[0].dtype == types.float32:
+            atomic_ref_ptr = builder.bitcast(
+                atomic_ref_ptr,
+                llvmir.PointerType(
+                    llvmir.IntType(32), addrspace=sig.args[0].address_space
+                ),
+            )
+            store_arg = builder.bitcast(store_arg, llvmir.IntType(32))
+        elif sig.args[0].dtype == types.float64:
+            atomic_ref_ptr = builder.bitcast(
+                atomic_ref_ptr,
+                llvmir.PointerType(
+                    llvmir.IntType(64), addrspace=sig.args[0].address_space
+                ),
+            )
+            store_arg = builder.bitcast(store_arg, llvmir.IntType(64))
 
         atomic_store_fn_args = [
-            builder.extract_value(
-                args[0],
-                context.data_model_manager.lookup(
-                    atomic_ref_ty
-                ).get_field_position("ref"),
-            ),
+            atomic_ref_ptr,
             context.get_constant(
                 types.int32, get_scope(atomic_ref_ty.memory_scope)
             ),
             context.get_constant(
                 types.int32,
                 get_memory_semantics_mask(atomic_ref_ty.memory_order),
             ),
-            args[1],
+            store_arg,
         ]
 
-        builder.call(atomic_store_fn, atomic_store_fn_args)
+        builder.call(
+            get_or_insert_spv_atomic_store_fn(
+                context, builder.module, atomic_ref_ty
+            ),
+            atomic_store_fn_args,
+        )
 
     return sig, _intrinsic_store_gen
 
@@ -298,28 +344,56 @@ def _intrinsic_exchange(
 
     def _intrinsic_exchange_gen(context, builder, sig, args):
         atomic_ref_ty = sig.args[0]
-        atomic_exchange_fn = get_or_insert_spv_atomic_exchange_fn(
-            context, builder.module, atomic_ref_ty
+
+        exchange_arg = args[1]
+        atomic_ref_ptr = builder.extract_value(
+            args[0],
+            context.data_model_manager.lookup(atomic_ref_ty).get_field_position(
+                "ref"
+            ),
         )
+        if sig.args[0].dtype == types.float32:
+            atomic_ref_ptr = builder.bitcast(
+                atomic_ref_ptr,
+                llvmir.PointerType(
+                    llvmir.IntType(32), addrspace=sig.args[0].address_space
+                ),
+            )
+            exchange_arg = builder.bitcast(exchange_arg, llvmir.IntType(32))
+        elif sig.args[0].dtype == types.float64:
+            atomic_ref_ptr = builder.bitcast(
+                atomic_ref_ptr,
+                llvmir.PointerType(
+                    llvmir.IntType(64), addrspace=sig.args[0].address_space
+                ),
+            )
+            exchange_arg = builder.bitcast(exchange_arg, llvmir.IntType(64))
 
         atomic_exchange_fn_args = [
-            builder.extract_value(
-                args[0],
-                context.data_model_manager.lookup(
-                    atomic_ref_ty
-                ).get_field_position("ref"),
-            ),
+            atomic_ref_ptr,
             context.get_constant(
                 types.int32, get_scope(atomic_ref_ty.memory_scope)
             ),
             context.get_constant(
                 types.int32,
                 get_memory_semantics_mask(atomic_ref_ty.memory_order),
             ),
-            args[1],
+            exchange_arg,
         ]
 
-        return builder.call(atomic_exchange_fn, atomic_exchange_fn_args)
+        ret_val = builder.call(
+            get_or_insert_spv_atomic_exchange_fn(
+                context, builder.module, atomic_ref_ty
+            ),
+            atomic_exchange_fn_args,
+        )
+
+        if sig.args[0].dtype == types.float32:
+            ret_val = builder.bitcast(ret_val, llvmir.FloatType())
+        elif sig.args[0].dtype == types.float64:
+            ret_val = builder.bitcast(ret_val, llvmir.DoubleType())
+
+        return ret_val
 
     return sig, _intrinsic_exchange_gen
 

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_group_barrier_overloads.py

@@ -6,7 +6,6 @@
 Provides overloads for functions included in kernel_iface.barrier that
 generate dpcpp SPIR-V LLVM IR intrinsic function calls.
 """
-import warnings
 
 from llvmlite import ir as llvmir
 from numba.core import cgutils, types
@@ -20,18 +19,7 @@
 from numba_dpex.kernel_api.memory_enums import MemoryOrder, MemoryScope
 
 from ._spv_atomic_inst_helper import get_memory_semantics_mask, get_scope
-
-_SUPPORT_CONVERGENT = True
-
-try:
-    llvmir.FunctionAttributes("convergent")
-except ValueError:
-    warnings.warn(
-        "convergent attribute is supported only starting llvmlite "
-        + "0.42. Not setting this attribute may result in unexpected behavior"
-        + "when using group_barrier"
-    )
-    _SUPPORT_CONVERGENT = False
+from .spv_atomic_fn_declarations import _SUPPORT_CONVERGENT
 
 
 def _get_memory_scope(fence_scope):

numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/spv_atomic_fn_declarations.py

@@ -7,22 +7,43 @@
 functions and their use inside an LLVM module.
 """
 
+import warnings
+
 from llvmlite import ir as llvmir
 from numba.core import cgutils, types
 
 from numba_dpex.core import itanium_mangler as ext_itanium_mangler
 from numba_dpex.kernel_api_impl.spirv.target import CC_SPIR_FUNC
 
+_SUPPORT_CONVERGENT = True
+
+try:
+    llvmir.FunctionAttributes("convergent")
+except ValueError:
+    warnings.warn(
+        "convergent attribute is supported only starting llvmlite "
+        + "0.42. Not setting this attribute may result in unexpected behavior"
+        + "when using group_barrier"
+    )
+    _SUPPORT_CONVERGENT = False
+
 
 def get_or_insert_atomic_load_fn(context, module, atomic_ref_ty):
     """
     Gets or inserts a declaration for a __spirv_AtomicLoad call into the
     specified LLVM IR module.
     """
     atomic_ref_dtype = atomic_ref_ty.dtype
+
+    if atomic_ref_dtype == types.float32:
+        atomic_ref_dtype = types.uint32
+    elif atomic_ref_dtype == types.float64:
+        atomic_ref_dtype = types.uint64
+
     atomic_load_fn_retty = context.get_value_type(atomic_ref_dtype)
     ptr_type = atomic_load_fn_retty.as_pointer()
     ptr_type.addrspace = atomic_ref_ty.address_space
+
     atomic_load_fn_arg_types = [
         ptr_type,
         llvmir.IntType(32),
@@ -44,6 +65,10 @@ def get_or_insert_atomic_load_fn(context, module, atomic_ref_ty):
     )
     fn.calling_convention = CC_SPIR_FUNC
 
+    if _SUPPORT_CONVERGENT:
+        fn.attributes.add("convergent")
+    fn.attributes.add("nounwind")
+
     return fn
 
 
@@ -53,9 +78,16 @@ def get_or_insert_spv_atomic_store_fn(context, module, atomic_ref_ty):
     specified LLVM IR module.
     """
     atomic_ref_dtype = atomic_ref_ty.dtype
+
+    if atomic_ref_dtype == types.float32:
+        atomic_ref_dtype = types.uint32
+    elif atomic_ref_dtype == types.float64:
+        atomic_ref_dtype = types.uint64
+
     ptr_type = context.get_value_type(atomic_ref_dtype).as_pointer()
     ptr_type.addrspace = atomic_ref_ty.address_space
     atomic_store_fn_retty = llvmir.VoidType()
+
     atomic_store_fn_arg_types = [
         ptr_type,
         llvmir.IntType(32),
@@ -80,6 +112,10 @@ def get_or_insert_spv_atomic_store_fn(context, module, atomic_ref_ty):
     )
     fn.calling_convention = CC_SPIR_FUNC
 
+    if _SUPPORT_CONVERGENT:
+        fn.attributes.add("convergent")
+    fn.attributes.add("nounwind")
+
     return fn
 
 
@@ -89,9 +125,16 @@ def get_or_insert_spv_atomic_exchange_fn(context, module, atomic_ref_ty):
     specified LLVM IR module.
     """
     atomic_ref_dtype = atomic_ref_ty.dtype
+
+    if atomic_ref_dtype == types.float32:
+        atomic_ref_dtype = types.uint32
+    elif atomic_ref_dtype == types.float64:
+        atomic_ref_dtype = types.uint64
+
     ptr_type = context.get_value_type(atomic_ref_dtype).as_pointer()
     ptr_type.addrspace = atomic_ref_ty.address_space
-    atomic_exchange_fn_retty = context.get_value_type(atomic_ref_ty.dtype)
+    atomic_exchange_fn_retty = context.get_value_type(atomic_ref_dtype)
+
     atomic_exchange_fn_arg_types = [
         ptr_type,
         llvmir.IntType(32),
@@ -118,6 +161,10 @@ def get_or_insert_spv_atomic_exchange_fn(context, module, atomic_ref_ty):
     )
     fn.calling_convention = CC_SPIR_FUNC
 
+    if _SUPPORT_CONVERGENT:
+        fn.attributes.add("convergent")
+    fn.attributes.add("nounwind")
+
     return fn
 
 
@@ -174,4 +221,8 @@ def get_or_insert_spv_atomic_compare_exchange_fn(
     )
     fn.calling_convention = CC_SPIR_FUNC
 
+    if _SUPPORT_CONVERGENT:
+        fn.attributes.add("convergent")
+    fn.attributes.add("nounwind")
+
     return fn