[mlir][linalg] Let tile and fuse fail for tile sizes zero.

[gysit]

Summary

Adapt tileConsumerAndFuseProducers to return failure if the generated tile loop nest is empty since all tile sizes are zero. Additionally, fix LinalgTileAndFuseTensorOpsPattern to return success if the pattern applied successfully.

[gysit]

This is a follow up to https://reviews.llvm.org/D118576. I am happy to keep the early exit if all tile sizes are zero just removed it to demonstrate that tileConsumerAndFuseProducers now returns failure if all tile sizes are zero (instead of an invalid tile loop nest).

[mravishankar]

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

595

Change looks fine. But would rather keep this to provide better diagnostics with -debug.

[gysit]

Address comment and fix test.

[gysit]

Fix warning.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

public:
   LinalgOp getRootOp() { return rootOp; }
 
   /// Returns the tiled root operation and the fused producers.
   SmallVector<LinalgOp> getAllTiledAndFusedOps();
 
   /// Returns the loop ops generated from tiling.
   ArrayRef<scf::ForOp> getLoopOps() { return tileLoopOps; }
 
-private:
   /// Returns true if the tile loop nest has no tile loops.
   bool isEmpty();
 
+private:
   /// Returns true if the tile loop nest invariants are satisfied:
   /// - The `rootOp` has been tiled at least once.
   /// - The number of tile loop operations and dimensions match.
   /// - The innermost tile loop is the parent of `tiledOp`.
   /// - The tile loops are directly nested.
   // TODO: relax to support additional control flow, e.g., IfOp.
   bool isValid();
 

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

mlir::linalg::tileConsumerAndFuseProducers(OpBuilder &b, LinalgOp consumerOp,
   // Tile the remaining loops and fuse the input operands.
   SmallVector<int64_t> innerTileSizes;
   innerTileSizes.append(split, 0);
   innerTileSizes.append(tileSizes.begin() + split, tileSizes.end());
   if (failed(tileLoopNest.tileRootOp(b, innerTileSizes, tileInterchange)))
     return failure();
   fuseProducersGreedily(tileLoopNest.getRootOp().getInputOperands());
 
+  // Exit if the tile loop nest is empty since all tile sizes are zero.
+  if (tileLoopNest.isEmpty())
+    return failure();
+
   return tileLoopNest;
 }

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

if (!options.tileInterchange.empty() &&
       options.tileInterchange.size() != options.tileSizes.size())
     return rewriter.notifyMatchFailure(
         op, "expect the number of tile sizes and interchange dims to match");
 
   // Copy the `tileSizes` and `tileInterchange` prefixes needed for `rootOp`.
   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()))
           : SmallVector<int64_t>(options.tileInterchange.begin(),
                                  options.tileInterchange.begin() +
                                      rootOp.getNumLoops());
 
+  // Check `rootTileSizes` contains non-zero tile sizes.
+  if (llvm::count(rootTileSizes, 0) == static_cast<long>(rootTileSizes.size()))
+    return rewriter.notifyMatchFailure(
+        op, "expect at least one non-zero tile size");
+
   // Check `rootInterchange` is a permutation of the `rootOp` loop dimensions.
   // It has to be a permutation since the tiling cannot tile the same loop
   // dimension multiple times.
   if (!isPermutation(rootInterchange))
     return rewriter.notifyMatchFailure(
         op, "expect the tile interchange permutes the root loops");
 
   // Tile `rootOp` and fuse its producers.
   FailureOr<TileLoopNest> tileLoopNest = tileConsumerAndFuseProducers(
       rewriter, rootOp, rootTileSizes, rootInterchange);
   if (failed(tileLoopNest))
     return rewriter.notifyMatchFailure(
         op, "tileConsumerAndFuseProducers failed unexpectedly");
 
   // Replace all uses of the tiled loop operation.
   rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults());
 
   // Apply the filter if specified.
   for (LinalgOp linalgOp : tileLoopNest->getAllTiledAndFusedOps())
     filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
-  return failure();
+  return success();
 }
 
 /// Linalg generic interchange pattern.
 mlir::linalg::GenericOpInterchangePattern::GenericOpInterchangePattern(
     MLIRContext *context, ArrayRef<unsigned> interchangeVector,
     LinalgTransformationFilter f, PatternBenefit benefit)
     : OpRewritePattern(context, benefit), filter(std::move(f)),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}

mlir/test/Dialect/Linalg/tile-and-fuse-no-fuse.mlir

-// 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
+// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.matmul fuse tile-sizes=0,0,0 run-enable-pass=false" -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>