Review comment v1.0

Author: Abhishek-Varma

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/KernelDispatch.cpp

@@ -303,6 +303,21 @@ FailureOr<ParameterSetting> ParameterSetting::create(
     TileSize maxL0Size = selectL2TileSizes(tileParams, m0Pack, n0Pack);
     M0 = maxL0Size.M;
     N0 = maxL0Size.N;
+
+    // TODO(avarma): This is currently a workaround for 1x1 AIE array to make
+    // those 2D matmul shapes work for which all of the operands get pulled in
+    // to L2 buffer. Once reprogramming of DMA ops is supported, we can get rid
+    // of this workaround. We need to add this only for pack-peel-4-level-tiling
+    // NOT pack-peel. The workaround just ensures that the tile size of first
+    // level is NOT equal to M, N by halving the N0 tile.
+    if (numRows == 1 && numCols == 1) {
+      // Check if the tile size generated is exactly same as operand size. If
+      // yes, halve N0 tile.
+      if (M0 == M && N0 == N) {
+        N0 /= 2;
+        if (N0 < n0Pack) n0Pack /= 2;
+      }
+    }
   }
 
   // Currently there is only one level of tiling for K dimension, and the packed
@@ -431,15 +446,6 @@ static LogicalResult setRootConfigForPackPeel4LevelTilingPipeline(
   SmallVector<SmallVector<int64_t>> outerPerm;
   SmallVector<PackingConfigPackingLevelAttr> packingConfigLevelsVal;
 
-  int64_t m0Tile = packPeelTiling.M0;
-  int64_t n0Tile = packPeelTiling.N0;
-  // For 4D matmul-like ops, only tile the outer dims.
-  // outer_tile_size = total_tile_size / inner_dim_size
-  if (is4DMatmulLikeOp(linalgOp)) {
-    m0Tile /= maybeInputDimsAndSizes.value().mSizes.back();
-    n0Tile /= maybeInputDimsAndSizes.value().nSizes.back();
-  }
-
   // Pack level => 1.
   // For 2D matmul-like ops, the first level is to pack operands from 2D to 4D.
   // If the input is a 4D matmul-like op, this level of packing is not needed.
@@ -450,30 +456,6 @@ static LogicalResult setRootConfigForPackPeel4LevelTilingPipeline(
     packedSizesL0[nDims.back()] = packPeelTiling.n0Pack;
     packedSizesL0[kDims.back()] = packPeelTiling.k0Pack;
 
-    // TODO(avarma): This is currently a workaround for 1x1 AIE array to make
-    // those 2D matmul shapes work for which all of the operands get pulled in
-    // to L2 buffer. Once reprogramming of DMA ops is supported, we can get rid
-    // of this workaround. We need to add this only for pack-peel-4-level-tiling
-    // NOT pack-peel. The workaround just ensures that the tile size of first
-    // level is NOT equal to M,N by halving the n0Tile.
-    if (numRows == 1 && numCols == 1) {
-      auto getTotalSize = [](ArrayRef<int64_t> sizes) {
-        return std::accumulate(sizes.begin(), sizes.end(), 1,
-                               std::multiplies<int64_t>());
-      };
-
-      // Get the shape (M, N) of the full Matmul operation.
-      auto maybeInputDimsAndSizes = getInputDimsAndSizes(linalgOp);
-      int64_t M = getTotalSize(maybeInputDimsAndSizes.value().mSizes);
-      int64_t N = getTotalSize(maybeInputDimsAndSizes.value().nSizes);
-      // Check if the tile size generated is exactly same as operand size. If
-      // yes, halve n0Tile.
-      if (m0Tile == M && n0Tile == N) {
-        n0Tile /= 2;
-        if (n0Tile < packedSizesL0[nDims.back()])
-          packedSizesL0[nDims.back()] /= 2;
-      }
-    }
     transposePackIndices = {0, 1, 2};
     // There is no corresponding unpack for the specified pack operation
     // 0 is used when unpack is empty
@@ -540,6 +522,14 @@ static LogicalResult setRootConfigForPackPeel4LevelTilingPipeline(
     assert(!batchDims.empty() && "expected batch dims not empty");
     tileSizeLevel0[batchDims[0]] = 1;
   }
+  int64_t m0Tile = packPeelTiling.M0;
+  int64_t n0Tile = packPeelTiling.N0;
+  // For 4D matmul-like ops, only tile the outer dims.
+  // outer_tile_size = total_tile_size / inner_dim_size
+  if (is4DMatmulLikeOp(linalgOp)) {
+    m0Tile /= maybeInputDimsAndSizes.value().mSizes.back();
+    n0Tile /= maybeInputDimsAndSizes.value().nSizes.back();
+  }
   tileSizeLevel0[mDims[0]] = m0Tile;
   tileSizeLevel0[nDims[0]] = n0Tile;
 

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/lowering_strategy_objectfifo_npu4.mlir

@@ -45,7 +45,7 @@ module {
 // -----
 
 // Tests a matmul shape for pack-peel-4-level-tiling in which the tile size generated
-// should ideally be equal to the M,N,K dimensions of the matmul - but it won't work
+// should ideally be equal to the M, N, K dimension size of the matmul - but it won't work
 // until support for DMA ops' reconfiguration is added. The workaround therefore halves
 // the tile size for N as N/2.
 
@@ -74,7 +74,7 @@ func.func @matmul_32x512x64_i32() {
 // -----
 
 // Based on above workaround this test shows the packing size of N also being halved
-// in case the tile size for N dimension becomes lesser than the corresponding packing size.
+// in case the tile size for N dimension becomes less than the corresponding packing size.
 
 // PACK-PEEL-4-LEVEL-1-CORE{LITERAL}: #config = #iree_codegen.lowering_config<tile_sizes = [[32, 16, 0], [1, 1, 0], [0, 0, 1], [1, 1, 0]]>
 // PACK-PEEL-4-LEVEL-1-CORE{LITERAL}: #packingConfig = #amdaie.packing_config<packing_config = [{packedSizes = [32, 16, 64], transposePackIndices = [0, 1, 2], unpackEmpty = [false, false, true], innerPerm = [[0, 1], [1, 0], [0, 1]], outerPerm = [[0, 1], [1, 0], [1, 0]]}, {packedSizes = [0, 0, 0, 4, 4, 8], transposePackIndices = [0, 1, 2], unpackEmpty = [false, false, true], innerPerm = [[0, 1], [1, 0], [0, 1]], outerPerm = [[0, 1, 3, 2], [0, 1, 3, 2], [0, 1, 3, 2]]}]>