[DEVICE] Adding ability to choose unroll factor at runtime (#1734)

* Adding runtime unroll factor selection via RCCL_UNROLL_FACTOR * [BUILD] Add support for user-defined UNROLL for debugging * Update CHANGELOG.md * Fix COLLTRACE errors in CI * Add debug statements for unroll and resolve warnings * Incorporate UNROLL into ONLY_FUNCS for debugging --------- Signed-off-by: nileshnegi <Nilesh.Negi@amd.com> Co-authored-by: gilbertlee-amd <44450918+gilbertlee-amd@users.noreply.github.com> Co-authored-by: Jeffrey Novotny <jnovotny@amd.com> [ROCm/rccl commit: 9d72be7b2f]
2025-06-11 00:07:59 -05:00
parent 4cd71722f2
commit 4cadf3597c
@@ -13,31 +13,6 @@ __shared__ ncclShmemData ncclShmem;
  __shared__ ulong2 ncclShmemPerWarp[ncclShmemScratchWarpSize()*(NCCL_MAX_NTHREADS/WARP_SIZE)/sizeof(ulong2)];
 #endif

-struct RunWorkNop {
-  __device__ void run() {}
-};
-
-__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
-  ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/1>(&args4K.args);
-}
-__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
-  ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/2>(&args4K.args);
-}
-__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernel_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
-  ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/false, /*Unroll*/4>(&args4K.args);
-}
-#ifdef ENABLE_COLLTRACE
-__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
-  ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/1>(&args4K.args);
-}
-__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
-  ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/2>(&args4K.args);
-}
-__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {
-  ncclKernelMain<-1, RunWorkNop, /*COLLTRACE*/true, /*Unroll*/4>(&args4K.args);
-}
-#endif
-
 #ifdef USE_INDIRECT_FUNCTION_CALL
 __device__ void ncclDevFunc_Nop();
 #else
@@ -11,8 +11,8 @@
 #include "collectives.h"
 #include "device.h"
 #include "op128.h"
-#include "reduce_kernel.h"
 #include "device_table.h"
+#include "reduce_kernel.h"
 #include "network/unpack/unpack_defs.h"
 #define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1)  // Using balanced tree instead of split tree

@@ -572,21 +572,7 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
    if (0 <= SpecializedFnId && ncclShmem.funcId == (unsigned)SpecializedFnId) {
      SpecializedRunWorkBatch().run();
    } else {
-#ifdef USE_INDIRECT_FUNCTION_CALL
-      if (COLL_UNROLL == 1)
-        ncclDevFuncTable_1[ncclShmem.funcId]();
-      else if (COLL_UNROLL == 2)
-        ncclDevFuncTable_2[ncclShmem.funcId]();
-      else
-        ncclDevFuncTable_4[ncclShmem.funcId]();
-#else
-      if (COLL_UNROLL == 1)
-        NCCL_CALL_FUNCTIONS_1(ncclShmem.funcId);
-      else if (COLL_UNROLL == 2)
-        NCCL_CALL_FUNCTIONS_2(ncclShmem.funcId);
-      else
-        NCCL_CALL_FUNCTIONS_4(ncclShmem.funcId);
-#endif
+      NCCL_CALL_FUNCTIONS<COLL_UNROLL>(ncclShmem.funcId);
    }

    if (ncclShmem.nextBatchIx == -1) break;
@@ -628,15 +614,6 @@ __device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* a
 #endif
 }

-__global__ void ncclDevKernel_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
-__global__ void ncclDevKernel_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
-__global__ void ncclDevKernel_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
-#ifdef ENABLE_COLLTRACE
-__global__ void ncclDevKernelDebug_Generic_1(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
-__global__ void ncclDevKernelDebug_Generic_2(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
-__global__ void ncclDevKernelDebug_Generic_4(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K);
-#endif
-
 #define DEFINE_ncclDevKernel_nop(suffix, coll, redop, ty, algo, proto, specializedFnId) \
  __global__ void ncclDevKernel_##suffix(ncclDevKernelArgs4K NCCL_GRID_CONSTANT const args4K) {}

@@ -32,7 +32,7 @@ else:
 # developing device code. The regex supports non-space containing globs '*',
 # and union 'a|b'. The string representing the function has the form:
 #
-# <coll> <algo> <proto> <redop> <type>
+# <coll> <algo> <proto> <redop> <type> <unroll>
 #
 # The possible values for redop, type, algo, proto can be found in the all_<foo>
 # lists at the top of this file.
@@ -45,23 +45,28 @@ else:
 # # Only AllReduce and Reduce
 # make ONLY_FUNCS="AllReduce|Reduce"
 #
+# # Only AllGather with unroll=4
+# make ONLY_FUNCS="AllGather * * * * 4"
+#
 # # Only non-reductions:
 # make ONLY_FUNCS="AllGather * *|Broadcast * *|SendRecv"
 #
 # # Only AllReduce Sum int32_t (but all algos, protos)
 # make ONLY_FUNCS="AllReduce * * Sum int32_t"
 #
-# # Only AllReduce RING Max float (but all protos)
+# # Only AllReduce RING Max float (but all protos and unrolls)
 # make ONLY_FUNCS="AllReduce RING * Max float"
 #
-# # AllReduce TREE LL128 Prod rccl_bfloat16
-# make ONLY_FUNCS="AllReduce TREE LL128 Prod rccl_bfloat16"
+# # AllReduce TREE LL128 Prod rccl_bfloat16 unroll=1
+# make ONLY_FUNCS="AllReduce TREE LL128 Prod rccl_bfloat16 1"
 #
-# # AllReduce RING SIMPLE and ReduceScatter RING LL float (but all redops, types for AllReduce and all redops for ReduceScatter)
-# make ONLY_FUNCS="AllReduce RING SIMPLE * *|ReduceScatter RING LL * float"
+# # AllReduce RING SIMPLE and ReduceScatter RING LL float (but all redops, types, unrolls for AllReduce and all redops, unrolls for ReduceScatter)
+# make ONLY_FUNCS="AllReduce RING SIMPLE * *|ReduceScatter RING LL * float *"
 #                         --- or ---
-# make ONLY_FUNCS="AllReduce RING SIMPLE|ReduceScatter RING LL * float"
-# make ONLY_FUNCS="AllReduce RING/TREE LL/SIMPLE Sum/MinMax int8_t/uint8_t/half/float/double/hip_bfloat16/rccl_float8/rccl_bfloat8|AllGather RING LL/SIMPLE Sum int8_t|AllToAllPivot RING SIMPLE Sum int8_t|Broadcast RING LL/SIMPLE Sum int8_t|Reduce RING LL/SIMPLE Sum/MinMax int8_t/uint8_t/half/float/double/hip_bfloat16/rccl_float8/rccl_bfloat8|ReduceScatter RING LL/SIMPLE Sum/MinMax int8_t/uint8_t/half/float/double/hip_bfloat16/rccl_float8/rccl_bfloat8|SendRecv RING SIMPLE Sum int8_t"
+# make ONLY_FUNCS="AllReduce RING SIMPLE|ReduceScatter RING LL * float *"
+# make ONLY_FUNCS="AllReduce RING/TREE LL/SIMPLE Sum/MinMax int8_t/uint8_t/half/float/double/hip_bfloat16/rccl_float8/rccl_bfloat8 1/2/4|AllGather RING LL/SIMPLE Sum int8_t 1/2/4|AllToAllPivot RING SIMPLE Sum int8_t 1/2/4|Broadcast RING LL/SIMPLE Sum int8_t 1/2/4|Reduce RING LL/SIMPLE Sum/MinMax int8_t/uint8_t/half/float/double/hip_bfloat16/rccl_float8/rccl_bfloat8 1/2/4|ReduceScatter RING LL/SIMPLE Sum/MinMax int8_t/uint8_t/half/float/double/hip_bfloat16/rccl_float8/rccl_bfloat8 1/2/4|SendRecv RING SIMPLE Sum int8_t 1/2/4"
+#
+# # ONLY_FUNCS can be used together for debugging

 # Paste all non-None arguments together with `sep`.
 def paste(sep, *args):
@@ -70,9 +75,8 @@ def paste(sep, *args):
 is_ifc             = 1 if sys.argv[2] == "ON" else 0
 is_colltrace       = 1 if sys.argv[3] == "ON" else 0
 is_msccl_kernels   = 1 if sys.argv[4] == "ON" else 0
-is_local_arch_only = 1 if sys.argv[5] == "ON" else 0

-func_pattern = sys.argv[6:7]
+func_pattern = sys.argv[5:]
 if func_pattern and func_pattern[0]:
  func_pattern = func_pattern[0]
 else:
@@ -133,49 +137,13 @@ coll_lower_to_camel = {coll_camel_to_lower[x]: x for x in coll_camel_to_lower}

 ################################################################################

-def calc_unroll_for_local_arch():
-  if not is_local_arch_only: 
-    return
-
-  rocminfo_path = os.environ.get('ROCM_PATH') + "/bin/rocminfo"
-
-  res = subprocess.run([rocminfo_path], stdout=subprocess.PIPE, universal_newlines=True)
-  rocminfo_output = res.stdout
-  
-  # Parse rocminfo binary output
-  gfx_targets = {}
-  curr_name = None
-  for line in rocminfo_output.splitlines():
-    line = line.strip()
-
-    if line.startswith("Name:"):
-      name = line.split(':')[-1].strip()
-      if "gfx" in name:
-        curr_name = name
-    if line.startswith("Compute Unit:") and curr_name:
-      cu_count = int(line.split(':')[-1].strip())
-      gfx_targets[(curr_name, cu_count)] = None
-      curr_name = None
-  
-  # We want to remove duplicates but cannot use a dictionary since same gfx name can have different cu counts
-  # Use (gfx_name, cu_count) as key for dictionary and convert it to list here
-  gfx_targets = list(gfx_targets.keys())
-
-  # Homogeneous system is required to build for only 1 varient of unroll factor
-  if len(gfx_targets) == 1:
-    gfx_name, cu_count = gfx_targets[0]
-    if "gfx950" == gfx_name:
-        return 1
-    elif "gfx908" == gfx_name or ("gfx942" == gfx_name and cu_count > 80):
-      return 2
-    else:
-      return 4
+seen_unroll = []

 # Helper function to check if the conditions for the collective is being met
-def func_validate(coll, algo, proto, redop, ty):
+def func_validate(coll, algo, proto, redop, ty, unroll):
  if redop == "SumPostDiv" and ty[0] not in ("i","u"):
    return False
-  if algo not in algos_of_coll[coll] or proto not in protos_of_coll[coll] or redop not in redops_of_coll[coll] or ty not in tys_of_coll[coll]:
+  if algo not in algos_of_coll[coll] or proto not in protos_of_coll[coll] or redop not in redops_of_coll[coll] or ty not in tys_of_coll[coll] or unroll not in all_unroll:
    return False
  return True

@@ -193,8 +161,8 @@ def func_filter(function_params, current_idx, item_list=None):
    if current_element == "*":
      if current_idx == 0:
        raise ValueError("Error: Paramter 'COLL' can not be type all '*'.")
-      
-      # all_params list must be in the same order as function_params --> <coll> <algo> <proto> <redop> <type>
+
+      # all_params list must be in the same order as function_params --> <coll> <algo> <proto> <redop> <type> <unroll>
      # Get the current list from all_params
      current_list = all_params[current_idx]

@@ -210,12 +178,12 @@ def func_filter(function_params, current_idx, item_list=None):
      # Check if the current element is recognized
      elements = current_element.split("/")
      current_param = all_params[current_idx]
-      
+
      # Iterate over the elements in the elements list
      for item in elements:
        if item not in current_param:
          raise ValueError(f"Error: {item} is unrecognized or does not belong to this category {current_param}.")
-        
+
      for item in elements:
        item_list.append(item)
        yield from func_filter(function_params, current_idx+1, item_list)
@@ -225,7 +193,9 @@ def func_filter(function_params, current_idx, item_list=None):
  else:
    coll, algo, proto, redop, ty, unroll = item_list

-    if func_validate(coll, algo, proto, redop, ty):
+    if func_validate(coll, algo, proto, redop, ty, unroll):
+      if not unroll in seen_unroll:
+        seen_unroll.append(unroll)
      yield(coll, algo, proto, redop, ty, unroll)

 # Parse ONLY_FUNCS input and feed it to func_filter
@@ -247,10 +217,6 @@ def parse_input(func_pattern):
 # Maps functions to the chosen representative for the equivalence class it
 # belongs to. For instance (sum, signed int) maps to (sum, unsigned int).
 def equivalent_primary(coll, algo, proto, redop, ty, unroll):
-  # if local arch only, we only need to build for 1 varient of coll_unroll.
-  # map the other varient of coll_unroll to this one.
-  if coll_unroll:
-    unroll = str(coll_unroll)
  if coll in ("AllReduce", "Reduce", "ReduceScatter"):
    # map signed integer sum/prod to unsigned
    if redop in ("Sum","Prod","PreMulSum","SumPostDiv") and ty[0]=="i":
@@ -268,13 +234,13 @@ def enumerate_func_rows():
        for proto in all_protos:
          for redop in all_redops:
            for ty in all_tys:
-              if func_validate(coll, algo, proto, redop, ty):
+              if func_validate(coll, algo, proto, redop, ty, unroll):
                yield (coll, algo, proto, redop, ty, unroll)

-# Sort the hashmap based on custom key <coll> <algo> <proto> <redop> <ty>
+# Sort the hashmap based on custom key <coll> <algo> <proto> <redop> <ty> <unroll>
 def custom_sort_key(fn):
    coll, algo, proto, redop, ty, unroll = fn
-    
+
    return (
        all_unroll.index(unroll),
        all_colls.index(coll),
@@ -286,8 +252,6 @@ def custom_sort_key(fn):

 ################################################################################

-coll_unroll = calc_unroll_for_local_arch()
-
 # Corresponds to ncclDevFuncRowToId[]
 func_rows = [fn for fn in enumerate_func_rows()]

@@ -304,6 +268,8 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f:
  print("-- Generating %s" % os.path.join(gensrc, "device_table.h"))
  out = f.write

+  out("#include \"common.h\"\n\n")
+
  if is_ifc: func_declaration = "__device__ void"
  else: func_declaration = "__device__ __attribute__((noinline)) void"

@@ -320,113 +286,86 @@ with open(os.path.join(gensrc, "device_table.h"), "w") as f:
  out("\n")

  out("typedef void(*ncclDevFuncPtr_t)();\n\n")
-  out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable_1[] = {\n")
-  index1 = 0
-  for fn in primary_funcs:
-    coll, algo, proto, redop, ty, unroll = fn
-    if unroll != "1": continue
-    sym = paste("_", "ncclDevFunc", *fn)
-    if fn[2] == "LL128":
-      out("#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) && defined(ENABLE_LL128)\n")
-      out("/*%4d*/ %s,\n#else\n" % (index1, sym))
-      fn_ll = fn[:2] + ("LL",) + fn[3:]
-      sym_ll = paste("_", "ncclDevFunc", *fn_ll)
-      out("/*%4d*/ %s,\n#endif\n" % (index1, sym_ll))
-    else:
-      out("/*%4d*/ %s,\n" % (index1, sym))
-    index1 += 1
-  out("nullptr};\n")
-  out("\n")
-  out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable_2[] = {\n")
-  index2 = 0
-  for fn in primary_funcs:
-    coll, algo, proto, redop, ty, unroll = fn
-    if unroll != "2": continue
-    sym = paste("_", "ncclDevFunc", *fn)
-    if fn[2] == "LL128":
-      out("#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) && defined(ENABLE_LL128)\n")
-      out("/*%4d*/ %s,\n#else\n" % (index2, sym))
-      fn_ll = fn[:2] + ("LL",) + fn[3:]
-      sym_ll = paste("_", "ncclDevFunc", *fn_ll)
-      out("/*%4d*/ %s,\n#endif\n" % (index2, sym_ll))
-    else:
-      out("/*%4d*/ %s,\n" % (index2, sym))
-    index2 += 1
-  out("nullptr};\n")
-  out("\n")
-  out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable_4[] = {\n")
-  index4 = 0
-  for fn in primary_funcs:
-    coll, algo, proto, redop, ty, unroll = fn
-    if unroll != "4": continue
-    sym = paste("_", "ncclDevFunc", *fn)
-    if fn[2] == "LL128":
-      out("#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) && defined(ENABLE_LL128)\n")
-      out("/*%4d*/ %s,\n#else\n" % (index4, sym))
-      fn_ll = fn[:2] + ("LL",) + fn[3:]
-      sym_ll = paste("_", "ncclDevFunc", *fn_ll)
-      out("/*%4d*/ %s,\n#endif\n" % (index4, sym_ll))
-    else:
-      out("/*%4d*/ %s,\n" % (index4, sym))
-    index4 += 1
-  out("nullptr};\n")
-  out("\n")
-  
+
+  # Generate function tables per unroll factor
+  tableIdx = 0
+  for curr_unroll in seen_unroll:
+    out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable_%s[] = {\n" % curr_unroll)
+    tableIdx = 0
+    for fn in primary_funcs:
+      coll, algo, proto, redop, ty, unroll = fn
+      if curr_unroll != unroll: continue
+      sym = paste("_", "ncclDevFunc", *fn)
+      if fn[2] == "LL128":
+        out("#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) && defined(ENABLE_LL128)\n")
+        out("/*%4d*/ %s,\n#else\n" % (tableIdx, sym))
+        fn_ll = fn[:2] + ("LL",) + fn[3:]
+        sym_ll = paste("_", "ncclDevFunc", *fn_ll)
+        out("/*%4d*/ %s,\n#endif\n" % (tableIdx, sym_ll))
+      else:
+        out("/*%4d*/ %s,\n" % (tableIdx, sym))
+      tableIdx += 1
+    out("nullptr};\n")
+    out("\n")
+
+  # Construct indirection function workaround
  if not is_ifc:
-    out("template<unsigned short f, unsigned short l>\n"
-      "struct Caller1 {\n"
+    out("template<int unroll, unsigned short f, unsigned short l>\n"
+      "struct Caller {\n"
      "  static __forceinline__ __device__ __host__\n"
-      "  void call1(unsigned short funcIndex) noexcept\n"
+      "  void call(unsigned short funcIndex) noexcept\n"
      "  {\n"
      "    constexpr unsigned short m = f + (l - f) / 2;\n"
-      "    return (funcIndex < m) ? Caller1<f, m>::call1(funcIndex) : Caller1<m, l>::call1(funcIndex);\n"
+      "    return (funcIndex < m) ? Caller<unroll, f, m>::call(funcIndex) : Caller<unroll, m, l>::call(funcIndex);\n"
      "  }\n"
      "};\n"
-      "\n"
-      "template<unsigned short f>\n"
-      "struct Caller1<f, f + 1>{\n"
-      "  static __forceinline__ __device__ __host__\n"
-      "  void call1(unsigned short funcIndex) noexcept { ncclDevFuncTable_1[f](); }\n"
-      "};\n")
-    out("__forceinline__ __device__ void NCCL_CALL_FUNCTIONS_1(unsigned short funcIndex) noexcept {\n")
-    out(f"  Caller1<0, {index1}>::call1(funcIndex);\n")
-    out("}\n\n")
-    out("template<unsigned short f, unsigned short l>\n"
-      "struct Caller2 {\n"
-      "  static __forceinline__ __device__ __host__\n"
-      "  void call2(unsigned short funcIndex) noexcept\n"
-      "  {\n"
-      "    constexpr unsigned short m = f + (l - f) / 2;\n"
-      "    return (funcIndex < m) ? Caller2<f, m>::call2(funcIndex) : Caller2<m, l>::call2(funcIndex);\n"
-      "  }\n"
-      "};\n"
-      "\n"
-      "template<unsigned short f>\n"
-      "struct Caller2<f, f + 1>{\n"
-      "  static __forceinline__ __device__ __host__\n"
-      "  void call2(unsigned short funcIndex) noexcept { ncclDevFuncTable_2[f](); }\n"
-      "};\n")
-    out("__forceinline__ __device__ void NCCL_CALL_FUNCTIONS_2(unsigned short funcIndex) noexcept {\n")
-    out(f"  Caller2<0, {index2}>::call2(funcIndex);\n")
-    out("}\n\n")
-    out("template<unsigned short f, unsigned short l>\n"
-      "struct Caller4 {\n"
-      "  static __forceinline__ __device__ __host__\n"
-      "  void call4(unsigned short funcIndex) noexcept\n"
-      "  {\n"
-      "    constexpr unsigned short m = f + (l - f) / 2;\n"
-      "    return (funcIndex < m) ? Caller4<f, m>::call4(funcIndex) : Caller4<m, l>::call4(funcIndex);\n"
-      "  }\n"
-      "};\n"
-      "\n"
-      "template<unsigned short f>\n"
-      "struct Caller4<f, f + 1>{\n"
-      "  static __forceinline__ __device__ __host__\n"
-      "  void call4(unsigned short funcIndex) noexcept { ncclDevFuncTable_4[f](); }\n"
-      "};\n")
-    out("__forceinline__ __device__ void NCCL_CALL_FUNCTIONS_4(unsigned short funcIndex) noexcept {\n")
-    out(f"  Caller4<0, {index4}>::call4(funcIndex);\n")
-    out("}\n\n")
+      "\n")
+
+    for curr_unroll in seen_unroll:
+      out("template<unsigned short f>\n")
+      out("struct Caller<%s, f, f + 1>{\n" % curr_unroll)
+      out("  static __forceinline__ __device__ __host__\n");
+      out("  void call(unsigned short funcIndex) noexcept { ncclDevFuncTable_%s[f](); }\n" % curr_unroll)
+      out("};\n")
+
+  out("\n")
+  # Create NCCL_CALL_FUNCTION helper function that will call the appropriate device function
+  out("template <int unroll>\n"
+      "__forceinline__ __device__ void NCCL_CALL_FUNCTIONS(unsigned short funcIndex) noexcept {\n")
+  if is_ifc:
+    for curr_unroll in seen_unroll:
+      out("  if (unroll == %s) { ncclDevFuncTable_%s[funcIndex]();\n" % (curr_unroll, curr_unroll))
+  else:
+    out(f"  Caller<unroll, 0, {tableIdx}>::call(funcIndex);\n")
+  out("}\n\n")
+
+  # Create RCCL
+  out("template<int SpecializedFnId, typename SpecializedRunWorkBatch, bool COLLTRACE, int COLL_UNROLL>\n");
+  out("__device__ __forceinline__ void ncclKernelMain(struct ncclDevKernelArgs const* args);\n\n");
+
+  out("struct RunWorkNop {\n");
+  out("  __device__ void run() {}\n");
+  out("};\n\n");
+
+  out("template <int UNROLL, bool COLLTRACE>\n"
+      "__launch_bounds__(NCCL_MAX_NTHREADS, 1) __global__ void rcclGenericKernel(ncclDevKernelArgs4K const args4K) {\n"
+      "  ncclKernelMain<-1, RunWorkNop, COLLTRACE, UNROLL>(&args4K.args);\n"
+      "}\n\n")
+
+  out("struct rcclKernelItem {\n");
+  out("  void* funcPtr;\n");
+  out("  int   unroll;\n");
+  out("};\n\n");
+
+  out("/* This table contains all the __global__ functions that were compiled */\n");
+  out("static struct rcclKernelItem rcclKernelTable[] = {\n")
+  for unroll in seen_unroll:
+    out("  {(void*)&(rcclGenericKernel<%s, false>), %s},\n" % (unroll, unroll))
+  out("#ifdef ENABLE_COLLTRACE\n")
+  for unroll in seen_unroll:
+    out("  {(void*)&(rcclGenericKernel<%s, true>), %s},\n" % (unroll, unroll))
+  out("#endif\n");
+  out("};\n\n");

 # Generate <gensrc>/device_table.cpp
 if is_colltrace:
@@ -436,7 +375,7 @@ if is_colltrace:
    out = f.write
    out('#include "nccl_common.h"\n#include "device.h"\n')
    out("\n")
-    
+
    seen_fns = set()
    out("const char* funcNames[FUNC_INDEX_TOTAL] = {\n")
    for fn in primary_funcs:
@@ -459,10 +398,13 @@ with open(os.path.join(gensrc, "host_table.cpp"), "w") as f:
  # The mapping from function rows to valid primary function ids.
  out("extern int const ncclDevFuncRowToId[] = {\n")
  index = 0
-  for fn in func_rows[:len(func_rows)//3]:
+  offset = len(func_rows)//len(all_unroll)
+  start = all_unroll.index(seen_unroll[0]) * offset
+  end = start + offset
+  for fn in func_rows[start:end]:
    fn_id, comment = -1, ""
    if fn is not None:
-      fn_id = primary_to_index[equivalent_primary(*fn)]
+      fn_id = primary_to_index[equivalent_primary(*fn)] % offset if primary_to_index[equivalent_primary(*fn)] != -1 else -1
      comment = " // " + paste(" ", *fn[:-1])
    out("/*%4d*/ %d,%s\n" % (index, fn_id, comment))
    index += 1
@@ -48,7 +48,10 @@
 } while (0)

 #define barrier_by_group() barrier_by_group_common(__threadfence())
+
+#if defined(__gfx942__) || defined(__gfx950__)
 #define barrier_by_group_block() barrier_by_group_common(__threadfence_block())
+#endif

 /* Protocol classes: ProtoSimple, ProtoLL, ProtoLL128
 * We use these as template args to the Primtiives class instead of integral