Check LHS batch pack

iree-org · Aug 29, 2023 · ab75999 · ab75999
1 parent 3726959
commit ab75999
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -1236,6 +1236,10 @@ getDefaultDistributionTileSizes(TilingInterface op) {
 }
 
 static bool isPackMatmulLHS(tensor::PackOp op) {
+  if (op.getSourceRank() == 3 && op.getInnerDimsPos().size() == 2 &&
+      op.getInnerDimsPos()[0] == 1 && op.getInnerDimsPos()[1] == 2) {
+    return true;
+  }
   return op.getSourceRank() == 2 && op.getInnerDimsPos().size() == 2 &&
          op.getInnerDimsPos()[0] == 0 && op.getInnerDimsPos()[1] == 1;
 }