Allow all matmul variants to be a root op (#5065)

iree-org · Mar 12, 2021 · e5cfa53 · e5cfa53
1 parent a5a0f1d
commit e5cfa53
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensors.cpp b/iree/compiler/Dialect/Flow/Transforms/DispatchLinalgOnTensors.cpp
@@ -180,9 +180,15 @@ static SmallVector<Value, 4> convertToWorkload(OpBuilder &b, Location loc,
 ///   linalg.init_tensor operations.
 
 static bool isRootOp(Operation *op) {
+  if (auto contractionOp = dyn_cast<linalg::ContractionOpInterface>(op)) {
+    if (contractionOp.isRowMajorMatmul() ||
+        contractionOp.isColumnMajorMatmul() ||
+        contractionOp.isRowMajorBatchMatmul()) {
+      return true;
+    }
+  }
   return isa<linalg::ConvInputNHWCFilterHWCFOp,
-             linalg::DepthwiseConvInputNHWCFilterHWCOp, linalg::MatmulOp,
-             linalg::BatchMatmulOp>(op);
+             linalg::DepthwiseConvInputNHWCFilterHWCOp>(op);
 }
 
 static bool isAlwaysClonedIntoDispatchOp(Operation *op) {