[Mlir-commits] [mlir] c8eed8f - [mlir][linalg] Assert tile loop nest invariants in fusion.

[Tobias Gysi]

Author: Tobias Gysi
Date: 2021-09-21T14:20:57Z
New Revision: c8eed8f9a7e019382ab236c90ac1ee0f5d81c54f

URL: https://github.com/llvm/llvm-project/commit/c8eed8f9a7e019382ab236c90ac1ee0f5d81c54f
DIFF: https://github.com/llvm/llvm-project/commit/c8eed8f9a7e019382ab236c90ac1ee0f5d81c54f.diff

LOG: [mlir][linalg] Assert tile loop nest invariants in fusion.

Assert the tile loop nest invariants are satisfied instead of failing silently.

Reviewed By: nicolasvasilache

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
index eedda7ddf7be4..0e65cb23964ca 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
@@ -301,8 +301,13 @@ LogicalResult TileLoopNest::tileRootOp(OpBuilder &b,
 
 FailureOr<LinalgOp> TileLoopNest::fuseProducer(OpBuilder &b,
                                                OpOperand *rootOpOperand) {
-  // Check the tile loop nest is non-empty and satisfies all invariants.
-  if (isEmpty() || !isValid())
+  assert(rootOpOperand->getOwner() == rootOp &&
+         "expect the root op to be the owner of the operand to fuse");
+  assert(this->isValid() &&
+         "expect the tile loop nest to satisfy all invariants");
+
+  // Check the tile loop nest is non-empty.
+  if (isEmpty())
     return failure();
 
   // Check `rootOpOperand` is defined by an ExtractSliceOp.
@@ -310,9 +315,8 @@ FailureOr<LinalgOp> TileLoopNest::fuseProducer(OpBuilder &b,
   if (!sliceOp)
     return failure();
 
-  // Check `tileLoopNest` tiles `sliceOp` and `rootOpOperand`.
-  if (sliceOp->getParentOp() != rootOp->getParentOp() ||
-      rootOpOperand->getOwner() != rootOp)
+  // Check `sliceOp` is tiled by the tile loop nest.
+  if (sliceOp->getParentOp() != rootOp->getParentOp())
     return failure();
 
   // Check if the producer is a LinalgOp possibly passed by iteration argument.