Add llvm.assume to spirv indexing functions

IntelPython · Mar 14, 2024 · fedd4d5 · fedd4d5
1 parent 2831862
commit fedd4d5
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_index_space_id_overloads.py b/numba_dpex/experimental/_kernel_dpcpp_spirv_overloads/_index_space_id_overloads.py
@@ -54,7 +54,7 @@ def _intrinsic_spirv_global_index_const(
     sig = types.int64(types.int32)
 
     def _intrinsic_spirv_global_index_const_gen(
-        context: SPIRVTargetContext,
+        context: SPIRVTargetContext,  # pylint: disable=unused-argument
         builder: llvmir.IRBuilder,
         sig,  # pylint: disable=unused-argument
         args,
@@ -79,7 +79,16 @@ def _intrinsic_spirv_global_index_const_gen(
             dim,
         )
 
-        return context.cast(builder, res, types.uintp, types.intp)
+        # Generating same check as sycl does. Did they add it to avoid pointer
+        # bitcast on special constant?
+        max_int32 = llvmir.Constant(res.type, 2147483648)
+        cmp = builder.icmp_unsigned("<", res, max_int32)
+
+        inst = builder.assume(cmp)
+        # TODO: tail does not always work
+        inst.tail = "tail"
+
+        return res
 
     return sig, _intrinsic_spirv_global_index_const_gen