[Mlir-commits] [mlir] 63baab8 - [mlir][sparse] Fold invariant op only when it has only one use.

[Peiming Liu]

Author: Peiming Liu
Date: 2022-10-28T23:42:28Z
New Revision: 63baab8b390e1c30fd7e2770fb40a99d0b05fe8d

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

LOG: [mlir][sparse] Fold invariant op only when it has only one use.

Reviewed By: aartbik

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 4946256d810c..4399caecadcc 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -169,7 +169,8 @@ struct FoldInvariantYield : public OpRewritePattern<GenericOp> {
   LogicalResult matchAndRewrite(GenericOp op,
                                 PatternRewriter &rewriter) const override {
     if (!op.hasTensorSemantics() || op.getNumResults() != 1 ||
-        !isAlloc(op.getDpsInitOperand(0), /*isZero=*/false) || !isZeroYield(op))
+        !isAlloc(op.getDpsInitOperand(0), /*isZero=*/false) ||
+        !isZeroYield(op) || !op.getDpsInitOperand(0)->get().hasOneUse())
       return failure();
     auto outputType = op.getResult(0).getType().cast<RankedTensorType>();
     // Yielding zero on newly allocated (all-zero) sparse tensors can be