Fix tile sizes

iree-org · Aug 25, 2023 · 4f0e6c6 · 4f0e6c6
1 parent dc19b29
commit 4f0e6c6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/Codegen/Interfaces/PartitionableLoopsInterface.cpp b/compiler/src/iree/compiler/Codegen/Interfaces/PartitionableLoopsInterface.cpp
@@ -83,6 +83,23 @@ struct Mmt4DOpPartitionableLoops
   }
 };
 
+/// External model implementation for linalg::BatchMmt4DOp.
+struct BatchMmt4DOpPartitionableLoops
+    : public PartitionableLoopsInterface::ExternalModel<
+          BatchMmt4DOpPartitionableLoops, linalg::BatchMmt4DOp> {
+  llvm::SmallVector<unsigned>
+  getPartitionableLoops(Operation *op,
+                        std::optional<unsigned> maxNumPartitionedLoops) const {
+    SmallVector<unsigned> parallelLoops = {0, 1, 2};
+    if (maxNumPartitionedLoops.has_value() &&
+        parallelLoops.size() > maxNumPartitionedLoops.value()) {
+      return llvm::to_vector(llvm::ArrayRef(parallelLoops)
+                                 .take_back(maxNumPartitionedLoops.value()));
+    }
+    return parallelLoops;
+  }
+};
+
 /// External model implementation for all operations to make only
 /// the outer parallel loops as partitionable.
 template <typename OpTy>
@@ -200,6 +217,11 @@ void registerInterfaceForLinalgOps<linalg::Mmt4DOp>(MLIRContext *ctx) {
   linalg::Mmt4DOp::attachInterface<Mmt4DOpPartitionableLoops>(*ctx);
 }
 
+template <>
+void registerInterfaceForLinalgOps<linalg::BatchMmt4DOp>(MLIRContext *ctx) {
+  linalg::BatchMmt4DOp::attachInterface<BatchMmt4DOpPartitionableLoops>(*ctx);
+}
+
 /// Registers the external models for all Linalg operations.
 template <typename OpTy1, typename OpTy2, typename... More>
 static void registerInterfaceForLinalgOps(MLIRContext *ctx) {

diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/test_config_mmt4d.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/test_config_mmt4d.mlir
@@ -90,7 +90,7 @@ hal.executable private @batch_mmt4d {
   }
 }
 
-//  CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[0, 0, 20, 0, 0, 0, 0], [1, 1, 1, 0, 8, 4, 0], [0, 0, 0, 1, 0, 0, 1]{{\]}}>
+//  CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 10, 20, 0, 0, 0, 0], [1, 1, 1, 0, 8, 4, 0], [0, 0, 0, 1, 0, 0, 1]{{\]}}>
 //      CHECK: func.func @batch_mmt4d()
 //      CHECK:   linalg.batch_mmt4d
 // CHECK-SAME:     lowering_config = #[[CONFIG]]