Revert "[BACKEND] Emit reqnctapercluster (#8645)"

This reverts commit 9b75018.

Author: whitneywhtsang

lib/Conversion/TritonGPUToLLVM/FuncOpToLLVM.cpp

@@ -179,25 +179,11 @@ struct FuncOpConversion : public ConvertOpToLLVMPattern<triton::FuncOp> {
             "ttg.total-num-warps"))
       numWarps = totalNumWarps.getInt();
 
-    int numCTAs = 1;
-    if (auto module = funcOp->getParentOfType<ModuleOp>()) {
-      if (auto moduleAttr =
-              module->getAttrOfType<IntegerAttr>(triton::gpu::AttrNumCTAsName))
-        numCTAs = moduleAttr.getInt();
-    }
-
     // Set `nvvm.maxnreg` if it was specified on the module.
     if (Attribute maxnregAttr =
             funcOp.getParentOp()->getAttr(triton::gpu::AttrMaxRegistersName))
       newFuncOp->setAttr(NVVM::NVVMDialect::getMaxnregAttrName(), maxnregAttr);
 
-    // Do we want to do this for nCTAs == 1 whenever sm >= 90?
-    if (numCTAs > 1) {
-      // Request a specific number of CTAs per cluster in the generated PTX.
-      newFuncOp->setAttr(NVVM::NVVMDialect::getClusterDimAttrName(),
-                         rewriter.getDenseI32ArrayAttr(numCTAs));
-    }
-
     // Set an attribute for reqntidx, it could be used in latter LLVM codegen
     // for `nvvm.annotation` metadata.
     newFuncOp->setAttr(NVVM::NVVMDialect::getReqntidAttrName(),

python/triton/compiler/compiler.py

@@ -294,6 +294,18 @@ def compile(src, target=None, options=None, _env_vars=None):
 
     metadata["cache_dir"] = fn_cache_manager.cache_dir
     metadata["triton_version"] = __version__
+    cluster_dims = getattr(options, "cluster_dims", None)
+    if cluster_dims is None:
+        num_ctas = getattr(options, "num_ctas", None)
+        if num_ctas is None:
+            num_ctas = 1
+        cluster_dims = (num_ctas, 1, 1)
+    if not isinstance(cluster_dims, (list, tuple)):
+        cluster_dims = (cluster_dims, )
+    cluster_dims = tuple(cluster_dims)
+    if len(cluster_dims) < 3:
+        cluster_dims = cluster_dims + (1, ) * (3 - len(cluster_dims))
+    metadata["cluster_dims"] = cluster_dims
     # run compilation pipeline  and populate metadata
     stages = dict()
     backend.add_stages(stages, options, src.language)
@@ -420,6 +432,7 @@ def __init__(self, src, metadata_group, hash):
         from collections import namedtuple
         metadata_path = next((Path(p) for c, p in metadata_group.items() if c.endswith(".json")))
         metadata = json.loads(metadata_path.read_text())
+        metadata['cluster_dims'] = tuple(metadata['cluster_dims'])
         # JSON serialization dumps the target as a dict. Restore it to a GPUTarget.
         target = metadata['target']
         metadata['target'] = GPUTarget(target['backend'], target['arch'], target['warp_size'])

test/Conversion/tritongpu_to_llvm_hopper.mlir

@@ -1,17 +1,5 @@
 // RUN: triton-opt %s -split-input-file --allocate-shared-memory-nv='compute-capability=90 ptx-version=81' --convert-triton-gpu-to-llvm='compute-capability=90 ptx-version=81' | FileCheck %s
 
-module attributes {"ttg.num-ctas" = 4 : i32, "ttg.num-warps" = 4 : i32} {
-  // CHECK-LABEL: @test_cluster_attr
-  // CHECK: nvvm.cluster_dim = array<i32: 4>
-  // CHECK: nvvm.kernel = 1 : ui1
-  // CHECK: nvvm.reqntid = array<i32: 128>
-  tt.func @test_cluster_attr(%lb : index, %A : !tt.ptr<f16>) {
-    tt.return
-  }
-}
-
-// -----
-
 #mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [8, 1], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0], instrShape = [16, 256, 32]}>
 #shared = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = false, elementBitWidth = 8}>
 #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 32, transposed = true, elementBitWidth = 8}>

third_party/amd/backend/compiler.py

@@ -34,6 +34,7 @@ class HIPOptions:
     num_stages: int = 2
     num_ctas: int = 1
     extern_libs: dict = None
+    cluster_dims: tuple = (1, 1, 1)
     debug: bool = False
     sanitize_overflow: bool = True
     arch: str = None
@@ -137,6 +138,9 @@ def pack_metadata(self, metadata):
             metadata.num_warps,
             metadata.num_ctas,
             metadata.shared,
+            metadata.cluster_dims[0],
+            metadata.cluster_dims[1],
+            metadata.cluster_dims[2],
         )
 
     def get_codegen_implementation(self, options):

third_party/amd/backend/driver.py

@@ -440,7 +440,7 @@ def format_of(ty):
 
 #define HIP_CHECK(ans) {{ gpuAssert((ans), __FILE__, __LINE__); }}
 
-static void _launch(int gridX, int gridY, int gridZ, int num_warps, int num_ctas, int launch_cooperative_grid, int shared_memory, hipStream_t stream, hipFunction_t function, hipDeviceptr_t profile_scratch{', ' + arg_decls if len(arg_decls) > 0 else ''}) {{
+static void _launch(int gridX, int gridY, int gridZ, int num_warps, int num_ctas, int launch_cooperative_grid, int clusterDimX, int clusterDimY, int clusterDimZ, int shared_memory, hipStream_t stream, hipFunction_t function, hipDeviceptr_t profile_scratch{', ' + arg_decls if len(arg_decls) > 0 else ''}) {{
   hipDeviceptr_t global_scratch = 0;
   void *params[] = {{ {', '.join(params)} }};
   if (gridX*gridY*gridZ > 0 && launch_cooperative_grid) {{
@@ -548,8 +548,8 @@ def format_of(ty):
   {' '.join(float_storage_decls)}
 
   // extract kernel metadata
-  int num_warps, num_ctas, shared_memory;
-  if (!PyArg_ParseTuple(kernel_metadata, \"iii\", &num_warps, &num_ctas, &shared_memory)) {{
+  int num_warps, num_ctas, shared_memory, clusterDimX, clusterDimY, clusterDimZ;
+  if (!PyArg_ParseTuple(kernel_metadata, \"iiiiii\", &num_warps, &num_ctas, &shared_memory, &clusterDimX, &clusterDimY, &clusterDimZ)) {{
     return NULL;
   }}
   // extract launch metadata
@@ -571,7 +571,7 @@ def format_of(ty):
 
   // raise exception asap
   {"; ".join([f"DevicePtrInfo ptr_info{i} = getPointer(_arg{i}, {i}); if (!ptr_info{i}.valid) return NULL;" if ty[0] == "*" else "" for i, ty in signature.items()])};
-  _launch(gridX, gridY, gridZ, num_warps, num_ctas, launch_cooperative_grid, shared_memory, (hipStream_t)_stream, (hipFunction_t)_function, (hipDeviceptr_t)profile_scratch{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''});
+  _launch(gridX, gridY, gridZ, num_warps, num_ctas, launch_cooperative_grid, clusterDimX, clusterDimY, clusterDimZ, shared_memory, (hipStream_t)_stream, (hipFunction_t)_function, (hipDeviceptr_t)profile_scratch{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''});
 
   if(launch_exit_hook != Py_None){{
     PyObject* ret = PyObject_CallOneArg(launch_exit_hook, launch_metadata);

third_party/nvidia/backend/compiler.py

@@ -204,6 +204,9 @@ def pack_metadata(self, metadata):
             metadata.num_warps,
             metadata.num_ctas,
             metadata.shared,
+            metadata.cluster_dims[0],
+            metadata.cluster_dims[1],
+            metadata.cluster_dims[2],
         )
 
     def get_codegen_implementation(self, options):
@@ -313,6 +316,8 @@ def make_ttgir(mod, metadata, opt, capability):
         passes.common.add_canonicalizer(pm)
 
         pm.run(mod, 'make_ttgir')
+        # num_ctas == 16 is non-portable. Does work for H100 and B200 tho
+        metadata["cluster_dims"] = (opt.num_ctas, 1, 1)
         metadata["tensordesc_meta"] = mod.get_tensordesc_metadata()
         return mod
 
@@ -331,6 +336,8 @@ def gluon_to_ttgir(self, src, metadata, options, capability):
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
 
         pm.run(mod, 'gluon_to_ttgir')
+        # num_ctas == 16 is non-portable. Does work for H100 and B200 tho
+        metadata["cluster_dims"] = (options.num_ctas, 1, 1)
         metadata["tensordesc_meta"] = mod.get_tensordesc_metadata()
         return mod
 

third_party/nvidia/backend/driver.c

@@ -231,11 +231,13 @@ defineGetFunctionHandle(getCuTensorMapEncodeTiledHandle,
                         cuTensorMapEncodeTiled);
 
 static PyObject *occupancyMaxActiveClusters(PyObject *self, PyObject *args) {
-  int clusterDim = -1, maxActiveClusters = -1;
+  int clusterDimX = -1, clusterDimY = -1, clusterDimZ = -1,
+      maxActiveClusters = -1;
   int shared = 0;
   CUfunction func;
 
-  if (!PyArg_ParseTuple(args, "Kii", &func, &shared, &clusterDim)) {
+  if (!PyArg_ParseTuple(args, "Kiiii", &func, &shared, &clusterDimX,
+                        &clusterDimY, &clusterDimZ)) {
     return NULL;
   }
 
@@ -248,13 +250,13 @@ static PyObject *occupancyMaxActiveClusters(PyObject *self, PyObject *args) {
 
   CUlaunchAttribute launchAttr[1];
   launchAttr[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
-  launchAttr[0].value.clusterDim.x = clusterDim;
-  launchAttr[0].value.clusterDim.y = 1;
-  launchAttr[0].value.clusterDim.z = 1;
+  launchAttr[0].value.clusterDim.x = clusterDimX;
+  launchAttr[0].value.clusterDim.y = clusterDimY;
+  launchAttr[0].value.clusterDim.z = clusterDimZ;
   CUlaunchConfig config;
-  config.gridDimX = clusterDim * maxActiveBlocks;
-  config.gridDimY = 1;
-  config.gridDimZ = 1;
+  config.gridDimX = clusterDimX;
+  config.gridDimY = maxActiveBlocks * clusterDimY;
+  config.gridDimZ = clusterDimZ;
   config.blockDimX = 128;
   config.blockDimY = 1;
   config.blockDimZ = 1;

third_party/nvidia/backend/driver.py

@@ -1,4 +1,5 @@
 import functools
+import operator
 import os
 import subprocess
 import triton
@@ -338,7 +339,7 @@ def format_of(ty):
 }}
 #endif
 
-static void _launch(int gridX, int gridY, int gridZ, int num_warps, int num_ctas, int launch_cooperative_grid, int launch_pdl, int shared_memory, CUstream stream, CUfunction function, CUdeviceptr global_scratch, CUdeviceptr profile_scratch{', ' + arg_decls if len(arg_decls) > 0 else ''}) {{
+static void _launch(int gridX, int gridY, int gridZ, int num_warps, int num_ctas, int launch_cooperative_grid, int launch_pdl, int clusterDimX, int clusterDimY, int clusterDimZ, int shared_memory, CUstream stream, CUfunction function, CUdeviceptr global_scratch, CUdeviceptr profile_scratch{', ' + arg_decls if len(arg_decls) > 0 else ''}) {{
   void *params[] = {{ {', '.join(params)} }};
   if (gridX*gridY*gridZ > 0) {{
     // 4 attributes that we can currently pass maximum
@@ -348,10 +349,16 @@ def format_of(ty):
       cuLaunchKernelExHandle = getLaunchKernelExHandle();
     }}
     CUlaunchConfig config;
-    config.gridDimX = gridX * num_ctas;
+    config.gridDimX = gridX;
     config.gridDimY = gridY;
     config.gridDimZ = gridZ;
 
+    if (num_ctas != 1) {{
+      config.gridDimX *= clusterDimX;
+      config.gridDimY *= clusterDimY;
+      config.gridDimZ *= clusterDimZ;
+    }}
+
     config.blockDimX = 32 * num_warps;
     config.blockDimY = 1;
     config.blockDimZ = 1;
@@ -375,9 +382,9 @@ def format_of(ty):
     if (num_ctas != 1) {{
       CUlaunchAttribute clusterAttr = {{}};
       clusterAttr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
-      clusterAttr.value.clusterDim.x = num_ctas;
-      clusterAttr.value.clusterDim.y = 1;
-      clusterAttr.value.clusterDim.z = 1;
+      clusterAttr.value.clusterDim.x = clusterDimX;
+      clusterAttr.value.clusterDim.y = clusterDimY;
+      clusterAttr.value.clusterDim.z = clusterDimZ;
       launchAttr[num_attrs] = clusterAttr;
       ++num_attrs;
 
@@ -388,7 +395,6 @@ def format_of(ty):
       ++num_attrs;
     }}
 
-    // num_ctas == 16 is non-portable. Does work for H100 and B200 tho
     config.numAttrs = num_attrs;
     if (num_ctas == 16) {{
       CUDA_CHECK(cuFuncSetAttribute(
@@ -534,8 +540,8 @@ def format_of(ty):
     return NULL;
   }}
 
-  int num_warps, num_ctas, shared_memory;
-  if (!PyArg_ParseTuple(kernel_metadata, \"iii\", &num_warps, &num_ctas, &shared_memory)) {{
+  int num_warps, num_ctas, shared_memory, clusterDimX, clusterDimY, clusterDimZ;
+  if (!PyArg_ParseTuple(kernel_metadata, \"iiiiii\", &num_warps, &num_ctas, &shared_memory, &clusterDimX, &clusterDimY, &clusterDimZ)) {{
     PyErr_SetString(PyExc_TypeError, "kernel_metadata must be a tuple");
     return NULL;
   }}
@@ -571,7 +577,7 @@ def format_of(ty):
   {newline.join(tma_decls)}
   {newline.join(float_storage_decls)}
   Py_BEGIN_ALLOW_THREADS;
-  _launch(gridX, gridY, gridZ, num_warps, num_ctas, launch_cooperative_grid, launch_pdl, shared_memory, (CUstream)_stream, (CUfunction)_function, global_scratch, profile_scratch{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''});
+  _launch(gridX, gridY, gridZ, num_warps, num_ctas, launch_cooperative_grid, launch_pdl, clusterDimX, clusterDimY, clusterDimZ, shared_memory, (CUstream)_stream, (CUfunction)_function, global_scratch, profile_scratch{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''});
   Py_END_ALLOW_THREADS;
   if (PyErr_Occurred()) {{
     return NULL;
@@ -713,7 +719,7 @@ def __init__(self, src, metadata):
             libraries=libraries,
         )
 
-        self.num_ctas = getattr(metadata, "num_ctas", 1)
+        self.num_ctas = functools.reduce(operator.mul, metadata.cluster_dims, 1)
         self.launch = wrap_handle_tensordesc(mod.launch, signature, tensordesc_meta)
         self.global_scratch_size = metadata.global_scratch_size
         self.global_scratch_align = metadata.global_scratch_align

third_party/proton/tutorials/matmul.py

@@ -24,7 +24,7 @@ def metadata_fn(
     grid_x, grid_y, grid_z = unpack_grid(grid)
     num_warps = metadata.num_warps
     num_stages = metadata.num_stages
-    cluster_x, cluster_y, cluster_z = unpack_grid((metadata.num_ctas, ))
+    cluster_x, cluster_y, cluster_z = metadata.cluster_dims
     shared_memory = metadata.shared
     M, K = args["a_ptr"].shape
     K, N = args["b_ptr"].shape