fixing input dtypes for load, store, and exchange functions using bitcast;adding convergent, nounwind attributes

Author: adarshyoga

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