Turn on blocking of contractions by default

of contractions by default due to regressions. This PR fixes the regressions and turns them on again. Fixes are: 1) If the original code had multiple dynamic dimensions, where one or more of those are not blocked, those need to be accounted for in the `expand_shape` generated. Signed-off-by: MaheshRavishankar <[email protected]>
iree-org · Nov 16, 2024 · 4653dc5 · 4653dc5
1 parent e3b6cc3
commit 4653dc5
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/BlockDynamicDimensions.cpp b/compiler/src/iree/compiler/Codegen/Common/BlockDynamicDimensions.cpp
@@ -22,7 +22,7 @@ static llvm::cl::opt<bool> clEnableBlockedMatmuls(
     "iree-codegen-block-dynamic-dimensions-of-contractions",
     llvm::cl::desc("developer flag to gaurd blocking dynamic dimensions of "
                    "contraction-like ops"),
-    llvm::cl::Hidden, llvm::cl::init(false));
+    llvm::cl::Hidden, llvm::cl::init(true));
 
 namespace mlir::iree_compiler {
 
@@ -125,14 +125,15 @@ blockDynamicDimensionsOfValue(RewriterBase &rewriter,
   SmallVector<OpFoldResult> outputShape;
   SmallVector<ReassociationIndices> reassociation;
   Location loc = v.getLoc();
+  SmallVector<OpFoldResult> origShape = tensor::getMixedSizes(rewriter, loc, v);
 
-  for (auto [index, dim] : llvm::enumerate(tensorType.getShape())) {
+  for (auto [index, dim] : llvm::enumerate(origShape)) {
     reassociation.emplace_back(ReassociationIndices{});
 
     // Check if this needs division.
     if (!tensorType.isDynamicDim(index) || !divisibilityInfo.contains(index)) {
       reassociation.back().push_back(outputShape.size());
-      outputShape.push_back(rewriter.getIndexAttr(dim));
+      outputShape.push_back(dim);
       continue;
     }
 
@@ -142,9 +143,8 @@ blockDynamicDimensionsOfValue(RewriterBase &rewriter,
     uint64_t factor = currDivisibility.sdiv();
     AffineExpr s0 = rewriter.getAffineSymbolExpr(0);
     AffineExpr divExpr = s0.floorDiv(factor);
-    Value sourceDim = rewriter.create<tensor::DimOp>(loc, v, index).getResult();
     OpFoldResult newDynamicDim = affine::makeComposedFoldedAffineApply(
-        rewriter, loc, divExpr, ArrayRef<OpFoldResult>{sourceDim});
+        rewriter, loc, divExpr, ArrayRef<OpFoldResult>{dim});
     OpFoldResult newStaticDim = rewriter.getIndexAttr(factor);
 
     reassociation.back().push_back(outputShape.size());

diff --git a/compiler/src/iree/compiler/Codegen/Common/test/block_dynamic_dims.mlir b/compiler/src/iree/compiler/Codegen/Common/test/block_dynamic_dims.mlir
@@ -228,3 +228,28 @@ func.func @reshape_propagation_test(%rhs : tensor<2048x4096xf16>, %m : index)
 //  CHECK-SAME:       outs(%[[EMPTY]] :
 //       CHECK:   %[[COLLAPSED:.+]] = tensor.collapse_shape %[[TRUNC]]
 //       CHECK:   return %[[COLLAPSED]]
+
+// -----
+
+func.func @multiple_dynamic_dims(%arg0 : index, %arg1 : index) -> tensor<?x?x4096xf32> {
+  %0 = util.assume.int %arg0<umin = 0, umax = 1024, udiv = 16> : index
+  %lhs = tensor.empty(%arg1, %0) : tensor<?x?x2048xf32>
+  %rhs = tensor.empty(%arg1) : tensor<?x2048x4096xf32>
+  %init = tensor.empty(%arg1, %0) : tensor<?x?x4096xf32>
+  %matmul = linalg.batch_matmul ins(%lhs, %rhs : tensor<?x?x2048xf32>, tensor<?x2048x4096xf32>)
+      outs(%init : tensor<?x?x4096xf32>) -> tensor<?x?x4096xf32>
+  return %matmul : tensor<?x?x4096xf32>
+}
+// CHECK-LABEL: func @multiple_dynamic_dims(
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
+//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index)
+//   CHECK-DAG:   %[[ARG0_ASSUME:.+]] = util.assume.int %[[ARG0]]
+//   CHECK-DAG:   %[[RHS:.+]] = tensor.empty(%[[ARG1]]) : tensor<?x2048x4096xf32>
+//   CHECK-DAG:   %[[BLOCKED_M:.+]] = affine.apply affine_map<()[s0] -> (s0 floordiv 16)>()[%[[ARG0_ASSUME]]]
+//   CHECK-DAG:   %[[LHS:.+]] = tensor.empty(%[[ARG1]], %[[BLOCKED_M]]) : tensor<?x?x16x2048xf32>
+//   CHECK-DAG:   %[[INIT:.+]] = tensor.empty(%[[ARG1]], %[[BLOCKED_M]]) : tensor<?x?x16x4096xf32>
+//       CHECK:   %[[MATMUL:.+]] = linalg.generic
+//  CHECK-SAME:       ins(%[[LHS]], %[[RHS]]
+//  CHECK-SAME:       outs(%[[INIT]] :
+//       CHECK:   %[[COLLAPSE:.+]] = tensor.collapse_shape %[[MATMUL]]
+//       CHECK:   return %[[COLLAPSE]]