[SYCL][PTX][CUDA] Implicit global offset implementation (#1773)

This commit implements implicit global offset behavior for the kernels generated for the PI CUDA backend. This includes the following changes:

 * A new builtin `__builtin_ptx_implicit_offset` and intrinsic `int.nvvm.implicit.offset` for getting the global offset. For the `ptx-nvidiacl` this is used for implementing the `__spirv_GlobalOffset` builtin.
 * A new pass that iterates over the uses of the `int.nvvm.implicit.offset` intrinsic, replacing it with a new function parameter. It then moves up the call-tree, adjusting calls to functions with this new parameter by adding a similar parameter to callers without it and adding this parameter to the calls. An exception are entry points, which are instead cloned with the clone being given the new parameter and the original using an offset of `{0,0,0}` in all uses of the intrinsic or functions with the new parameter. Any entry points that are not cloned are invariant to the offset parameter.

Additionally the PI CUDA backend now includes an offset parameter in the set of arguments for kernels. PI CUDA attempt to load the corresponding kernel both with and without the global offset parameter. If present, the kernel with the offset parameter is used only when a non-zero global offset is used.

Author: steffenlarsen

clang/include/clang/Basic/BuiltinsNVPTX.def

@@ -89,6 +89,9 @@ BUILTIN(__nvvm_read_ptx_sreg_pm1, "i", "n")
 BUILTIN(__nvvm_read_ptx_sreg_pm2, "i", "n")
 BUILTIN(__nvvm_read_ptx_sreg_pm3, "i", "n")
 
+// SYCL
+BUILTIN(__builtin_ptx_implicit_offset, "Ui*", "nc")
+
 // MISC
 
 BUILTIN(__nvvm_prmt, "UiUiUiUi", "")

libclc/ptx-nvidiacl/libspirv/workitem/get_global_id.cl

@@ -9,13 +9,16 @@
 #include <spirv/spirv.h>
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_GlobalInvocationId_x() {
-  return __spirv_WorkgroupId_x() * __spirv_WorkgroupSize_x() + __spirv_LocalInvocationId_x();
+  return __spirv_WorkgroupId_x() * __spirv_WorkgroupSize_x() +
+         __spirv_LocalInvocationId_x() + __spirv_GlobalOffset_x();
 }
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_GlobalInvocationId_y() {
-  return __spirv_WorkgroupId_y() * __spirv_WorkgroupSize_y() + __spirv_LocalInvocationId_y();
+  return __spirv_WorkgroupId_y() * __spirv_WorkgroupSize_y() +
+         __spirv_LocalInvocationId_y() + __spirv_GlobalOffset_y();
 }
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_GlobalInvocationId_z() {
-  return __spirv_WorkgroupId_z() * __spirv_WorkgroupSize_z() + __spirv_LocalInvocationId_z();
+  return __spirv_WorkgroupId_z() * __spirv_WorkgroupSize_z() +
+         __spirv_LocalInvocationId_z() + __spirv_GlobalOffset_z();
 }

libclc/ptx-nvidiacl/libspirv/workitem/get_global_offset.cl

@@ -11,13 +11,13 @@
 // Compiler support is required to provide global offset on NVPTX.
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_GlobalOffset_x() {
-    return 0;
+  return __builtin_ptx_implicit_offset()[0];
 }
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_GlobalOffset_y() {
-    return 0;
+  return __builtin_ptx_implicit_offset()[1];
 }
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_GlobalOffset_z() {
-    return 0;
+  return __builtin_ptx_implicit_offset()[2];
 }

llvm/include/llvm/IR/IntrinsicsNVVM.td

@@ -4161,4 +4161,10 @@ foreach layout_a = ["row", "col"] in {
   } // layout_b
 } // layout_a
 
+// SYCL
+def int_nvvm_implicit_offset :
+  GCCBuiltin<"__builtin_ptx_implicit_offset">,
+  Intrinsic<[LLVMPointerType<llvm_i32_ty>], [],
+  [IntrNoMem, IntrSpeculatable]>;
+
 } // let TargetPrefix = "nvvm"

llvm/lib/Target/NVPTX/CMakeLists.txt

@@ -33,6 +33,7 @@ set(NVPTXCodeGen_sources
   NVVMIntrRange.cpp
   NVVMReflect.cpp
   NVPTXProxyRegErasure.cpp
+  SYCL/GlobalOffset.cpp
   SYCL/LocalAccessorToSharedMemory.cpp
   )
 

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -16,8 +16,9 @@
 #include "NVPTXLowerAggrCopies.h"
 #include "NVPTXTargetObjectFile.h"
 #include "NVPTXTargetTransformInfo.h"
-#include "TargetInfo/NVPTXTargetInfo.h"
+#include "SYCL/GlobalOffset.h"
 #include "SYCL/LocalAccessorToSharedMemory.h"
+#include "TargetInfo/NVPTXTargetInfo.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
@@ -71,6 +72,7 @@ void initializeNVPTXLowerArgsPass(PassRegistry &);
 void initializeNVPTXLowerAllocaPass(PassRegistry &);
 void initializeNVPTXProxyRegErasurePass(PassRegistry &);
 
+void initializeGlobalOffsetPass(PassRegistry &);
 void initializeLocalAccessorToSharedMemoryPass(PassRegistry &);
 
 } // end namespace llvm
@@ -94,6 +96,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   initializeNVPTXProxyRegErasurePass(PR);
 
   // SYCL-specific passes, needed here to be available to `opt`.
+  initializeGlobalOffsetPass(PR);
   initializeLocalAccessorToSharedMemoryPass(PR);
 }
 
@@ -274,6 +277,7 @@ void NVPTXPassConfig::addIRPasses() {
 
   if (getTM<NVPTXTargetMachine>().getTargetTriple().getOS() == Triple::CUDA &&
       getTM<NVPTXTargetMachine>().getTargetTriple().getEnvironment() == Triple::SYCLDevice) {
+    addPass(createGlobalOffsetPass());
     addPass(createLocalAccessorToSharedMemoryPass());
   }