[AMD] GetThreadsPerWarpForOperand interface (#5675)

This PR implements GetThreadsPerWarpForOperand function for WMMA and
MFMA layouts.

Author: binarman

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1964,9 +1964,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>
@@ -2032,9 +2053,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 {

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 {