Guard the blocking of dynamic dimensions of contractions. (#19164)

The change to block dynamic dimensions of contractions based on divisibility analysis seems to have regressed quite a few tests. The change is probably too aggressive. Guarding that by a flag while I fix those forward. Signed-off-by: MaheshRavishankar <[email protected]>
iree-org · Nov 15, 2024 · 43d3d1a · 43d3d1a
1 parent f750087
commit 43d3d1a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/BlockDynamicDimensions.cpp b/compiler/src/iree/compiler/Codegen/Common/BlockDynamicDimensions.cpp
@@ -18,6 +18,12 @@
 
 #define DEBUG_TYPE "iree-codegen-block-dynamic-dimensions"
 
+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));
+
 namespace mlir::iree_compiler {
 
 #define GEN_PASS_DEF_BLOCKDYNAMICDIMENSIONSPASS
@@ -290,7 +296,10 @@ blockDynamicDimensions(RewriterBase &rewriter,
                                       attentionOp);
       })
       .Case<linalg::LinalgOp>([&](auto linalgOp) {
-        return blockDynamicDimensions(rewriter, dynamicDimAnalysis, linalgOp);
+        if (clEnableBlockedMatmuls) {
+          return blockDynamicDimensions(rewriter, dynamicDimAnalysis, linalgOp);
+        }
+        return success();
       })
       .Default([&](Operation *op) { return success(); });
 }

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
@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-block-dynamic-dimensions, cse))" --split-input-file --mlir-print-local-scope %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-block-dynamic-dimensions, cse))" --split-input-file --mlir-print-local-scope --iree-codegen-block-dynamic-dimensions-of-contractions %s | FileCheck %s
 
 #pipeline_layout = #hal.pipeline.layout<constants = 4, bindings = [
     #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">,