Avoid doing tile + fuse if tile sizes are zero.
authorMahesh Ravishankar <ravishankarm@google.com>
Tue, 1 Feb 2022 16:54:05 +0000 (16:54 +0000)
committerMahesh Ravishankar <ravishankarm@google.com>
Tue, 1 Feb 2022 18:34:06 +0000 (18:34 +0000)
Reviewed By: gysit

Differential Revision: https://reviews.llvm.org/D118576

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
mlir/test/Dialect/Linalg/tile-and-fuse-no-fuse.mlir [new file with mode: 0644]

index 7dcc977..f986667 100644 (file)
@@ -592,6 +592,10 @@ LogicalResult mlir::linalg::LinalgTileAndFuseTensorOpsPattern::matchAndRewrite(
   SmallVector<int64_t> rootTileSizes(options.tileSizes.begin(),
                                      options.tileSizes.begin() +
                                          rootOp.getNumLoops());
+  if (llvm::all_of(rootTileSizes, [](int64_t ts) { return ts == 0; })) {
+    return rewriter.notifyMatchFailure(
+        op, "all tile sizes are zero, nothing to do");
+  }
   SmallVector<int64_t> rootInterchange =
       options.tileInterchange.empty()
           ? llvm::to_vector<6>(llvm::seq<int64_t>(0, rootOp.getNumLoops()))
diff --git a/mlir/test/Dialect/Linalg/tile-and-fuse-no-fuse.mlir b/mlir/test/Dialect/Linalg/tile-and-fuse-no-fuse.mlir
new file mode 100644 (file)
index 0000000..80fc9bb
--- /dev/null
@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.generic fuse tile-sizes=0,0 run-enable-pass=false" -cse -split-input-file | FileCheck %s
+
+builtin.func @no_fuse_gemm(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 0.0 : f32
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
+  %fill = linalg.fill(%cst, %init) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+  %result = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %result : tensor<?x?xf32>
+}
+// CHECK-LABEL: @no_fuse_gemm(
+//   CHECK-NOT:   scf.for
+//       CHECK:   linalg.fill
+//   CHECK-NOT:   scf.for
+//       CHECK:   linalg.matmul