[SYCL][Fusion] Support heterogeneous ND ranges on HIP (#11805)

This PR enables support for fusing kernels with heterogeneous ND ranges
for the HIP target.

This leverages the `TargetFusionInterface` established in #11421,
however, due to the lack of explicit intrinsics for querying the
global/workgroup sizes in the AMDGCN backend, we had to extend the
interface further to be able to remap *loads from the dispatch pointer
with special offsets* as builtins as well.

---------

Signed-off-by: Julian Oppermann <julian.oppermann@codeplay.com>
Co-authored-by: Victor Lomuller <victor@codeplay.com>

Author: jopperm

libclc/amdgcn-amdhsa/libspirv/workitem/get_local_size.cl

@@ -23,17 +23,19 @@
 CONST_AS char * __clc_amdgcn_dispatch_ptr(void) __asm("llvm.amdgcn.dispatch.ptr");
 #endif
 
+// Mimic `EmitAMDGPUWorkGroupSize` in `clang/lib/CodeGen/CGBuiltin.cpp`.
+
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_WorkgroupSize_x() {
-    CONST_AS uint * ptr = (CONST_AS uint *) __dispatch_ptr();
-    return ptr[1] & 0xffffu;
+    CONST_AS ushort * ptr = (CONST_AS ushort *) __dispatch_ptr();
+    return ptr[2];
 }
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_WorkgroupSize_y() {
-    CONST_AS uint * ptr = (CONST_AS uint *) __dispatch_ptr();
-    return ptr[1] >> 16;
+    CONST_AS ushort * ptr = (CONST_AS ushort *) __dispatch_ptr();
+    return ptr[3];
 }
 
 _CLC_DEF _CLC_OVERLOAD size_t __spirv_WorkgroupSize_z() {
-    CONST_AS uint * ptr = (CONST_AS uint *) __dispatch_ptr();
-    return ptr[2] & 0xffffu;
+    CONST_AS ushort * ptr = (CONST_AS ushort *) __dispatch_ptr();
+    return ptr[4];
 }

sycl-fusion/jit-compiler/lib/KernelFusion.cpp

@@ -100,12 +100,6 @@ FusionResult KernelFusion::fuseKernels(
         "Fusion output target format not supported by this build");
   }
 
-  if (TargetFormat != BinaryFormat::SPIRV &&
-      TargetFormat != BinaryFormat::PTX && IsHeterogeneousList) {
-    return FusionResult{
-        "Heterogeneous ND ranges not supported for this target"};
-  }
-
   bool CachingEnabled = ConfigHelper::get<option::JITEnableCaching>();
   CacheKeyT CacheKey{KernelsToFuse,
                      Identities,

sycl-fusion/passes/kernel-fusion/Builtins.cpp

@@ -36,10 +36,9 @@ static raw_ostream &operator<<(raw_ostream &Os, const NDRange &ND) {
             << ND.getLocalSize();
 }
 
-/// Will generate a unique function name so that it can be reused in further
-/// stages.
-static std::string getFunctionName(BuiltinKind K, const NDRange &SrcNDRange,
-                                   const NDRange &FusedNDRange) {
+std::string Remapper::getFunctionName(BuiltinKind K, const NDRange &SrcNDRange,
+                                      const NDRange &FusedNDRange,
+                                      uint32_t Idx) {
   std::string Res;
   raw_string_ostream S{Res};
   S << "__" <<
@@ -63,6 +62,8 @@ static std::string getFunctionName(BuiltinKind K, const NDRange &SrcNDRange,
         llvm_unreachable("Unhandled kind");
       }()
     << "_remapper_" << SrcNDRange << "_" << FusedNDRange;
+  if (Idx != (uint32_t)-1)
+    S << "_" << static_cast<char>('x' + Idx);
   return S.str();
 }
 
@@ -339,13 +340,8 @@ jit_compiler::Remapper::remapBuiltins(Function *F, const NDRange &SrcNDRange,
         // If the builtin should not be remapped, return the original function.
         return F;
 
-      // Remap given builtin.
-      const auto Name = getFunctionName(K, SrcNDRange, FusedNDRange);
-      auto *M = F->getParent();
-      assert(!M->getFunction(Name) && "Function name should be unique");
-
       return Cached = TargetInfo.createRemapperFunction(
-                 *this, K, F->getName(), Name, M, SrcNDRange, FusedNDRange);
+                 *this, K, F, F->getParent(), SrcNDRange, FusedNDRange);
     }
     if (TargetInfo.isSafeToNotRemapBuiltin(F)) {
       // No need to remap.
@@ -375,20 +371,9 @@ jit_compiler::Remapper::remapBuiltins(Function *F, const NDRange &SrcNDRange,
 
   // Set Cached to support recursive functions.
   Cached = Clone;
-  for (auto &I : instructions(Clone)) {
-    if (auto *Call = dyn_cast<CallBase>(&I)) {
-      // Recursive call
-      auto *OldF = Call->getCalledFunction();
-      auto ErrOrNewF = remapBuiltins(OldF, SrcNDRange, FusedNDRange);
-      if (auto Err = ErrOrNewF.takeError()) {
-        return std::move(Err);
-      }
-      // Override called function.
-      auto *NewF = *ErrOrNewF;
-      Call->setCalledFunction(NewF);
-      Call->setCallingConv(NewF->getCallingConv());
-      Call->setAttributes(NewF->getAttributes());
-    }
+  if (auto Err = TargetInfo.scanForBuiltinsToRemap(Clone, *this, SrcNDRange,
+                                                   FusedNDRange)) {
+    return Err;
   }
   return Clone;
 }

sycl-fusion/passes/kernel-fusion/Builtins.h

@@ -47,6 +47,12 @@ class Remapper {
   explicit Remapper(const llvm::TargetFusionInfo &TargetInfo)
       : TargetInfo(TargetInfo) {}
 
+  ///
+  /// Generate a unique function name for a remapper function.
+  static std::string getFunctionName(BuiltinKind K, const NDRange &SrcNDRange,
+                                     const NDRange &FusedNDRange,
+                                     uint32_t Idx = -1);
+
   ///
   /// Recursively remap index space getters builtins.
   llvm::Expected<llvm::Function *> remapBuiltins(llvm::Function *F,