[mlir][sparse] replace zero yield generic op with copy in allocation

[aartbik]

Summary

This prepares patterns that sometimes are generated by the front-end
and would prohibit fusion of SDDMM flavored kernels.

[Peiming]

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

120

Why only fold zero?

[aartbik]

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

120

For now, the only situation that appears from automatically generated code.

[springerm]

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

102

What if the op has multiple outputs?

102

In addition, you could also check if yield.getOperand(0) is a zero constant.

120

initial

120

converts

134–139

Is this safe? What if a already has a copy operand and a has multiple users?

138

Should be wrapped in rewriter.updateRootInPlace.

139

Can you use zero here? arith.constant bufferize to a memref.get_global and is "non-writable". I.e., if some operation tries to write to it, the bufferization will automatically insert an alloc_tensor op.

[aartbik]

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

102

That is taken care of in L125. Or do you mean you want an assert here?

102

Extended this for direct value. Like Peiming suggested, for now I restrict this to zero (since that is the only useful case in the set of rewriting, but we may expand other constants too in the future).

134–139

We guarantee there is no copy at L127 (isZero == false). And even if there are several uses, the contents will be filled either way, right?

138

I have changed it, also below. See if this is the right idiom?

139

But that is why we replace the copy in the already existing bufferization allocation, right?

[aartbik]

addressed comments

[aartbik]

more precise update root call

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

 //===- SparseTensorRewriting.cpp - Sparse tensor rewriting rules ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements rewriting rules that are specific to sparse tensors.
 //
 //===----------------------------------------------------------------------===//
 
+#include "CodegenUtils.h"
+
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
if (isa<arith::AddFOp>(def) || isa<arith::AddIOp>(def)) {
       Value x = op.getBlock()->getArguments().back();
       return (def->getOperand(0) == x && isMulChain(def->getOperand(1), x)) ||
              (def->getOperand(1) == x && isMulChain(def->getOperand(0), x));
     }
   }
   return false;
 }
 
+// Helper to detect direct yield of a zero value.
+static bool isZeroYield(GenericOp op) {
+  auto yieldOp = cast<linalg::YieldOp>(op.region().front().getTerminator());
+  if (auto arg = yieldOp.getOperand(0).dyn_cast<BlockArgument>()) {
+    if (arg.getOwner()->getParentOp() == op) {
+      OpOperand *t = op.getInputAndOutputOperands()[arg.getArgNumber()];
+      return matchPattern(t->get(), m_Zero()) ||
+             matchPattern(t->get(), m_AnyZeroFloat());
+    }
+  } else if (auto *def = yieldOp.getOperand(0).getDefiningOp()) {
+    return matchPattern(def, m_Zero()) || matchPattern(def, m_AnyZeroFloat());
+  }
+  return false;
+}
+
 //===---------------------------------------------------------------------===//
 // The actual sparse tensor rewriting rules.
 //===---------------------------------------------------------------------===//
 
 namespace {
 
+/// Rewriting rule that converts direct yield of zero with initial allocation.
+struct FoldInvariantYield : public OpRewritePattern<GenericOp> {
+public:
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!op.hasTensorSemantics() || op.getNumResults() != 1 ||
+        !isAlloc(op.getOutputOperand(0), /*isZero=*/false) || !isZeroYield(op))
+      return failure();
+    auto outputType = op.getResult(0).getType().cast<RankedTensorType>();
+    if (!outputType.hasStaticShape() || getSparseTensorEncoding(outputType))
+      return failure();
+    // Incorporate zero value into allocation copy.
+    Value zero = constantZero(rewriter, op.getLoc(), op.getResult(0).getType());
+    AllocTensorOp a =
+        op.getOutputOperand(0)->get().getDefiningOp<AllocTensorOp>();
+    rewriter.updateRootInPlace(a, [&]() { a.getCopyMutable().assign(zero); });
+    rewriter.replaceOp(op, op.getOutputOperand(0)->get());
+    return success();
+  }
+};
+
 /// Rewriting rule that converts two kernels:
 ///
 ///      T(i,j) = SUM(k, A(i,j,k) * B(i,j,k) * ... )
 ///      X(i,j) = S(i,j) * T(i,j)
 ///
 /// into a single kernel, using distributive law:
 ///
 ///      X(i,j) = SUM(k, S(i,j) * A(i,j,k) * B(i,j,k) * ... )
for (auto &op : prodBlock.without_terminator())
         rewriter.clone(op, mapper);
       }
     mapper.map(consBlock.getArgument(other), fusedBlock->back().getResult(0));
     mapper.map(last, rewriter.clone(*sampler, mapper)->getResult(0));
     last = rewriter.clone(*acc, mapper)->getResult(0);
     rewriter.create<linalg::YieldOp>(loc, last);
     // Force initial value on merged allocation for dense outputs.
     if (!getSparseTensorEncoding(op.getResult(0).getType())) {
-      AllocTensorOp a1 =
-          prod.getOutputOperand(0)->get().getDefiningOp<AllocTensorOp>();
-      AllocTensorOp a2 =
+      Value init = prod.getOutputOperand(0)
+                       ->get()
+                       .getDefiningOp<AllocTensorOp>()
+                       .getCopy();
+      AllocTensorOp a =
           op.getOutputOperand(0)->get().getDefiningOp<AllocTensorOp>();
-      a2.getCopyMutable().assign(a1.getCopy());
+      rewriter.updateRootInPlace(a, [&]() { a.getCopyMutable().assign(init); });
     }
     // Replace consumer with fused operation. Old producer
     // and consumer ops will be removed by DCE.
     rewriter.replaceOp(op, fusedOp->getResults());
     return success();
   }
 
 private:

 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 
 void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns) {
-  patterns
-      .add<FuseSparseMultiplyOverAdd, ReshapeRewriter<tensor::ExpandShapeOp>,
+  patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd,
+               ReshapeRewriter<tensor::ExpandShapeOp>,
                ReshapeRewriter<tensor::CollapseShapeOp>>(patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir

#trait_scale = {
   indexing_maps = [
     affine_map<(d0, d1) -> (d0, d1)>,
     affine_map<(d0, d1) -> (d0, d1)>,
     affine_map<(d0, d1) -> (d0, d1)>
   ],
   iterator_types = ["parallel", "parallel"]
 }
 
+// CHECK-LABEL: func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> {
+// CHECK:         %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf64>
+// CHECK:         %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<1024x1024xf64>
+// CHECK:         return %[[VAL_1]] : tensor<1024x1024xf64>
+// CHECK:       }
+func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> {
+  %cst = arith.constant 0.000000e+00 : f64
+  %0 = linalg.init_tensor [1024, 1024] : tensor<1024x1024xf64>
+  %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>,
+                                        affine_map<(d0, d1) -> (d0, d1)>],
+                                        iterator_types = ["parallel", "parallel"]}
+                                        ins(%cst : f64)
+                                        outs(%0 : tensor<1024x1024xf64>) {
+    ^bb0(%a: f64, %x: f64):
+      linalg.yield %a : f64
+    } -> tensor<1024x1024xf64>
+  return %1 : tensor<1024x1024xf64>
+}
+
+// CHECK-LABEL: func.func @fold_yield_direct_zero() -> tensor<32xf64> {
+// CHECK:         %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<32xf64>
+// CHECK:         %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<32xf64>
+// CHECK:         return %[[VAL_1]] : tensor<32xf64>
+// CHECK:       }
+func.func @fold_yield_direct_zero() -> tensor<32xf64> {
+  %cst = arith.constant 0.000000e+00 : f64
+  %0 = linalg.init_tensor [32] : tensor<32xf64>
+  %1 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>],
+                                        iterator_types = ["parallel"]}
+                                        outs(%0 : tensor<32xf64>) {
+    ^bb0(%x: f64):
+      linalg.yield %cst : f64
+    } -> tensor<32xf64>
+  return %1 : tensor<32xf64>
+}
+
 // CHECK-LABEL: func.func @sampled_dd_unfused(
 // CHECK-SAME:    %[[VAL_0:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
 // CHECK-SAME:    %[[VAL_1:.*]]: tensor<8x8xf64>,
 // CHECK-SAME:    %[[VAL_2:.*]]: tensor<8x8xf64>) -> tensor<8x8xf64> {
 // CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 8 : index
 // CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[VAL_6:.*]] = arith.constant dense<0.000000e+00> : tensor<8x8xf64>