Merge commit 'b1301d66d19ad2244d24861505e176ef58bf8609'

Author: whitneywhtsang

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -2009,9 +2009,30 @@ AMDMfmaEncodingAttr::getRepOrderForOperand(int opIdx) const {
 
 SmallVector<unsigned>
 AMDMfmaEncodingAttr::getThreadsPerWarpForOperand(int opIdx) const {
-  llvm::report_fatal_error(
-      "getThreadsPerWarpForOperand not implemented for AMDMfmaEncodingAttr");
-  return {};
+  auto rank = ::getOrder(*this).size();
+  SmallVector<unsigned> threads(rank, 1);
+  unsigned kThreads;
+  unsigned nonKThreads;
+  switch (getMDim()) {
+  case 32:
+    assert(getNDim() == 32);
+    kThreads = 2;
+    nonKThreads = 32;
+    break;
+  case 16:
+    assert(getNDim() == 16);
+    kThreads = 4;
+    nonKThreads = 16;
+    break;
+  default:
+    llvm::report_fatal_error(
+        "unexpected mfma shape encountered in getThreadsPerWarpForOperand");
+  }
+  int kDimIdx = opIdx == 0 ? rank - 1 : rank - 2;
+  int nonKDimIdx = opIdx == 0 ? rank - 2 : rank - 1;
+  threads[kDimIdx] = kThreads;
+  threads[nonKDimIdx] = nonKThreads;
+  return threads;
 }
 
 SmallVector<int64_t>
@@ -2077,9 +2098,30 @@ AMDWmmaEncodingAttr::getRepOrderForOperand(int opIdx) const {
 
 SmallVector<unsigned>
 AMDWmmaEncodingAttr::getThreadsPerWarpForOperand(int opIdx) const {
-  llvm::report_fatal_error("getThreadsPerWarpForOperand not implemented for "
-                           "AMDWmmaEncodingAttr");
-  return {};
+  auto rank = ::getOrder(*this).size();
+  SmallVector<unsigned> threads(rank, 1);
+  unsigned kThreads;
+  unsigned nonKThreads;
+  switch (getVersion()) {
+  case 1:
+    // kThreads * onKThreads != 32,
+    // because values in lanes (n, n + 16) duplicates
+    kThreads = 1;
+    nonKThreads = 16;
+    break;
+  case 2:
+    kThreads = 2;
+    nonKThreads = 16;
+    break;
+  default:
+    llvm::report_fatal_error(
+        "unsupported WMMA version in getThreadsPerWarpForOperand");
+  }
+  int kDimIdx = opIdx == 0 ? rank - 1 : rank - 2;
+  int nonKDimIdx = opIdx == 0 ? rank - 2 : rank - 1;
+  threads[kDimIdx] = kThreads;
+  threads[nonKDimIdx] = nonKThreads;
+  return threads;
 }
 
 SmallVector<unsigned> AMDWmmaEncodingAttr::getCTAsPerCGA() const {

python/test/unit/language/test_core.py

@@ -1773,7 +1773,6 @@ def kernel_r(ptrs, BLOCK_SIZE: tl.constexpr):
 
 
 @pytest.mark.interpreter
-@pytest.mark.skipif(is_hip(), reason="Not implemented for AMD At this moment")
 def test_load_scope_sem_coop_grid_cta_one(device):
 
     @triton.jit

third_party/amd/backend/compiler.py

@@ -52,8 +52,6 @@ class HIPOptions:
     default_dot_input_precision: str = "ieee"
     allowed_dot_input_precisions: Tuple[str] = ("ieee", )
     enable_fp_fusion: bool = True
-    # TODO: Implement cooperative grid launch for AMD:
-    # See: https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Driver_API_functions_supported_by_HIP.html
     launch_cooperative_grid: bool = False
     matrix_instr_nonkdim: int = 0
     kpack: int = 1

third_party/amd/backend/driver.py

@@ -225,7 +225,7 @@ def format_of(ty):
         }[ty_to_cpp(ty)]
 
     args_format = ''.join([format_of(ty) for ty in signature.values()])
-    format = "iiiKKOOOO" + args_format
+    format = "piiiKKOOOO" + args_format
     signature = ','.join(map(_serialize_signature, signature.values()))
     signature = list(filter(bool, signature.split(',')))
     signature = {i: s for i, s in enumerate(signature)}
@@ -267,6 +267,12 @@ def format_of(ty):
                   unsigned int blockDimY, unsigned int blockDimZ,             \\
                   unsigned int sharedMemBytes, hipStream_t stream,            \\
                   void **kernelParams, void **extra)                          \\
+  FOR_EACH_ERR_FN(hipModuleLaunchCooperativeKernel, hipFunction_t f,          \\
+                  unsigned int gridDimX, unsigned int gridDimY,               \\
+                  unsigned int gridDimZ, unsigned int blockDimX,              \\
+                  unsigned int blockDimY, unsigned int blockDimZ,             \\
+                  unsigned int sharedMemBytes, hipStream_t stream,            \\
+                  void **kernelParams, void **extra)                          \\
   FOR_EACH_ERR_FN(hipPointerGetAttribute, void *data,                         \\
                   hipPointer_attribute attribute, hipDeviceptr_t ptr)
 
@@ -338,14 +344,18 @@ 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 clusterDimX, int clusterDimY, int clusterDimZ, int shared_memory, hipStream_t stream, hipFunction_t function{', ' + 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{', ' + arg_decls if len(arg_decls) > 0 else ''}) {{
   // printf("_launch hip kernel\\n");
   hipDeviceptr_t global_scratch = 0;
   void *params[] = {{ {', '.join(params)} }};
+  if (gridX*gridY*gridZ > 0 && launch_cooperative_grid) {{
+    HIP_CHECK(hipSymbolTable.hipModuleLaunchCooperativeKernel(function, gridX, gridY, gridZ, {warp_size}*num_warps, 1, 1, shared_memory, stream, params, 0));
+    return;
+  }}
   if (gridX*gridY*gridZ > 0) {{
-      HIP_CHECK(hipSymbolTable.hipModuleLaunchKernel(function, gridX, gridY, gridZ, {warp_size}*num_warps, 1, 1, shared_memory, stream, params, 0));
-    }}
+    HIP_CHECK(hipSymbolTable.hipModuleLaunchKernel(function, gridX, gridY, gridZ, {warp_size}*num_warps, 1, 1, shared_memory, stream, params, 0));
   }}
+}}
 
 typedef struct _DevicePtrInfo {{
     hipDeviceptr_t dev_ptr;
@@ -398,12 +408,14 @@ def format_of(ty):
   int gridX, gridY, gridZ;
   uint64_t _stream;
   uint64_t _function;
+  int launch_cooperative_grid;
   PyObject *launch_enter_hook = NULL;
   PyObject *launch_exit_hook = NULL;
   PyObject *kernel_metadata = NULL;
   PyObject *launch_metadata = NULL;
   {' '.join([f"{_extracted_type(ty)} _arg{i}; " for i, ty in signature.items()])}
-  if(!PyArg_ParseTuple(args, \"{format}\", &gridX, &gridY, &gridZ, &_stream, &_function,
+  if(!PyArg_ParseTuple(args, \"{format}\", &launch_cooperative_grid,
+                                           &gridX, &gridY, &gridZ, &_stream, &_function,
                                            &kernel_metadata, &launch_metadata,
                                            &launch_enter_hook, &launch_exit_hook {args_list})) {{
     return NULL;
@@ -426,7 +438,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, clusterDimX, clusterDimY, clusterDimZ, shared_memory, (hipStream_t)_stream, (hipFunction_t)_function{', ' + ', '.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{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''});
 
   if(launch_exit_hook != Py_None){{
     PyObject* args = Py_BuildValue("(O)", launch_metadata);
@@ -482,9 +494,10 @@ def __init__(self, src, metadata):
         src = make_launcher(constants, signature, metadata.warp_size)
         mod = compile_module_from_src(src, "__triton_launcher")
         self.launch = mod.launch
+        self.launch_cooperative_grid = metadata.launch_cooperative_grid
 
     def __call__(self, *args):
-        self.launch(*args)
+        self.launch(self.launch_cooperative_grid, *args)
 
 
 class HIPDriver(GPUDriver):

unittest/Dialect/TritonGPU/DialectTest.cpp

@@ -368,6 +368,8 @@ TEST_F(AMDMfmaLayoutTest, mfma_dot_op) {
   auto dot2dOp1 = createDotOperand(1, mfma2d, 4);
   ASSERT_THAT(dot2dOp0.getWarpOrder(), mfma2d.getWarpOrder());
   ASSERT_THAT(dot2dOp1.getWarpOrder(), mfma2d.getWarpOrder());
+  ASSERT_THAT(dot2dOp0.getThreadsPerWarp(), testing::ElementsAre(32u, 2u));
+  ASSERT_THAT(dot2dOp1.getThreadsPerWarp(), testing::ElementsAre(2u, 32u));
 
   auto tmfma2d = createTransposedMFMA(32, 32, {2, 4});
   auto tdot2dOp0 = createDotOperand(0, tmfma2d, 4);
@@ -380,12 +382,28 @@ TEST_F(AMDMfmaLayoutTest, mfma_dot_op) {
   auto dot3dOp1 = createDotOperand(1, mfma3d, 4);
   ASSERT_THAT(dot3dOp0.getWarpOrder(), mfma3d.getWarpOrder());
   ASSERT_THAT(dot3dOp1.getWarpOrder(), mfma3d.getWarpOrder());
+  ASSERT_THAT(dot3dOp0.getThreadsPerWarp(), testing::ElementsAre(1u, 32u, 2u));
+  ASSERT_THAT(dot3dOp1.getThreadsPerWarp(), testing::ElementsAre(1u, 2u, 32u));
 
   auto tmfma3d = createTransposedMFMA(32, 32, {2, 4, 1});
   auto tdot3dOp0 = createDotOperand(0, tmfma3d, 4);
   auto tdot3dOp1 = createDotOperand(1, tmfma3d, 4);
   ASSERT_THAT(tdot3dOp0.getWarpOrder(), tmfma3d.getWarpOrder());
   ASSERT_THAT(tdot3dOp1.getWarpOrder(), tmfma3d.getWarpOrder());
+
+  auto mfma16_2d = createMFMA(16, 16, {2, 4});
+  auto dot16_2dOp0 = createDotOperand(0, mfma16_2d, 4);
+  auto dot16_2dOp1 = createDotOperand(1, mfma16_2d, 4);
+  ASSERT_THAT(dot16_2dOp0.getThreadsPerWarp(), testing::ElementsAre(16u, 4u));
+  ASSERT_THAT(dot16_2dOp1.getThreadsPerWarp(), testing::ElementsAre(4u, 16u));
+
+  auto mfma16_3d = createMFMA(16, 16, {2, 4, 1});
+  auto dot16_3dOp0 = createDotOperand(0, mfma16_3d, 4);
+  auto dot16_3dOp1 = createDotOperand(1, mfma16_3d, 4);
+  ASSERT_THAT(dot16_3dOp0.getThreadsPerWarp(),
+              testing::ElementsAre(1u, 16u, 4u));
+  ASSERT_THAT(dot16_3dOp1.getThreadsPerWarp(),
+              testing::ElementsAre(1u, 4u, 16u));
 }
 
 TEST_F(AMDWmmaLayoutTest, wmmaV1) {
@@ -434,24 +452,36 @@ TEST_F(AMDWmmaLayoutTest, wmma_dot_op) {
   auto dot2dVer1Op1 = createDotOperand(1, wmma2dVer1, 16);
   ASSERT_THAT(dot2dVer1Op0.getWarpOrder(), wmma2dVer1.getWarpOrder());
   ASSERT_THAT(dot2dVer1Op1.getWarpOrder(), wmma2dVer1.getWarpOrder());
+  ASSERT_THAT(dot2dVer1Op0.getThreadsPerWarp(), testing::ElementsAre(16u, 1u));
+  ASSERT_THAT(dot2dVer1Op1.getThreadsPerWarp(), testing::ElementsAre(1u, 16u));
 
-  auto wmma3dVer1 = createWMMAv1({2, 4});
+  auto wmma3dVer1 = createWMMAv1({2, 4, 1});
   auto dot3dVer1Op0 = createDotOperand(0, wmma3dVer1, 16);
   auto dot3dVer1Op1 = createDotOperand(1, wmma3dVer1, 16);
   ASSERT_THAT(dot3dVer1Op0.getWarpOrder(), wmma3dVer1.getWarpOrder());
   ASSERT_THAT(dot3dVer1Op1.getWarpOrder(), wmma3dVer1.getWarpOrder());
+  ASSERT_THAT(dot3dVer1Op0.getThreadsPerWarp(),
+              testing::ElementsAre(1, 16u, 1u));
+  ASSERT_THAT(dot3dVer1Op1.getThreadsPerWarp(),
+              testing::ElementsAre(1, 1u, 16u));
 
   auto wmma2dVer2 = createWMMAv2(false, {2, 4});
   auto dot2dVer2Op0 = createDotOperand(0, wmma2dVer2, 16);
   auto dot2dVer2Op1 = createDotOperand(1, wmma2dVer2, 16);
   ASSERT_THAT(dot2dVer2Op0.getWarpOrder(), wmma2dVer2.getWarpOrder());
   ASSERT_THAT(dot2dVer2Op1.getWarpOrder(), wmma2dVer2.getWarpOrder());
+  ASSERT_THAT(dot2dVer2Op0.getThreadsPerWarp(), testing::ElementsAre(16u, 2u));
+  ASSERT_THAT(dot2dVer2Op1.getThreadsPerWarp(), testing::ElementsAre(2u, 16u));
 
-  auto wmma3dVer2 = createWMMAv2(false, {2, 4});
+  auto wmma3dVer2 = createWMMAv2(false, {2, 4, 1});
   auto dot3dVer2Op0 = createDotOperand(0, wmma3dVer2, 16);
   auto dot3dVer2Op1 = createDotOperand(1, wmma3dVer2, 16);
   ASSERT_THAT(dot3dVer2Op0.getWarpOrder(), wmma3dVer2.getWarpOrder());
   ASSERT_THAT(dot3dVer2Op1.getWarpOrder(), wmma3dVer2.getWarpOrder());
+  ASSERT_THAT(dot3dVer2Op0.getThreadsPerWarp(),
+              testing::ElementsAre(1, 16u, 2u));
+  ASSERT_THAT(dot3dVer2Op1.getThreadsPerWarp(),
+              testing::ElementsAre(1, 2u, 16u));
 }
 
 class LinearEncodingTest : public ::testing::Test {