[mlir][sparse] preserve zero-initialization for materializing buffers

[aartbik]

Summary

This revision makes sure that when the output buffer materializes locally
(in contrast with the passing in of output tensors either in-place or not
in-place), the zero initialization assumption is preserved. This also adds
a bit more documentation on our sparse kernel assumption (viz. TACO
assumptions).

[aartbik]

rebase with main

[bixia]

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

396

s/any a/any/

[aartbik]

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

396

Note that it was any "a priori" but since this is confusing, I rephrased

[aartbik]

rephrased doc

[aartbik]

rebased with main

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

static bool getInPlace(Value val) {
   if (auto arg = val.dyn_cast<BlockArgument>())
     if (auto funcOp = dyn_cast<FuncOp>(arg.getOwner()->getParentOp()))
       if (auto attr = funcOp.getArgAttrOfType<BoolAttr>(
               arg.getArgNumber(), linalg::LinalgDialect::kInplaceableAttrName))
         return attr.getValue();
   return false;
 }
 
-/// Generates buffer for the output tensor.
+/// Generates buffer for the output tensor. Note that all sparse kernels
+/// assume that when all elements are written to (viz. x(i) = y(i) * z(i)),
+/// the output buffer is already initialized to all zeroes and only nonzeroes
+/// values are computed and written out. For updates (viz. x(i) += y(i) * z(i)),
+/// only nonzeroes values are used for the updates and no assumption on the
+/// original contents of the output buffer is necessary..
 static Value genOutputBuffer(CodeGen &codegen, PatternRewriter &rewriter,
                              linalg::GenericOp op, MemRefType denseTp,
                              ArrayRef<Value> args) {
   Location loc = op.getLoc();
   Value tensor = op.getOutputOperand(0)->get();
   // The output tensor simply could materialize from the buffer that will
   // be generated for the tensor present in the outs() clause. This has
   // the major advantage that the sparse kernel only updates the nonzero
   // positions for the output tensor.
   if (getInPlace(tensor))
     return rewriter.create<memref::BufferCastOp>(loc, denseTp, tensor);
   // By default, a new buffer is allocated which is initialized to the
   // tensor defined in the outs() clause. This is always correct but
   // introduces a dense initialization component that may negatively
-  // impact the running complexity of the sparse kernel.
+  // impact the running complexity of the sparse kernel. If the tensor
+  // materializes within this method, we need to preserve the zero
+  // initialization assumption of all sparse output buffers.
+  if (auto init = tensor.getDefiningOp<linalg::InitTensorOp>()) {
+    Type tp = denseTp.getElementType();
+    Value alloc = rewriter.create<memref::AllocOp>(loc, denseTp, args);
+    Value zero = rewriter.create<ConstantOp>(loc, tp, rewriter.getZeroAttr(tp));
+    rewriter.create<linalg::FillOp>(loc, zero, alloc);
+    return alloc;
+  }
   Value init = rewriter.create<memref::BufferCastOp>(loc, denseTp, tensor);
   Value alloc = rewriter.create<memref::AllocOp>(loc, denseTp, args);
   rewriter.create<memref::CopyOp>(loc, init, alloc);
   return alloc;
 }
 
 /// Local bufferization of all dense and sparse data structures.
 /// This code enables testing the first prototype sparse compiler.

mlir/test/Dialect/SparseTensor/sparse_1d.mlir

%0 = linalg.generic #trait1
     outs(%argx: tensor<32xf32>) {
       ^bb(%a: f32, %x: f32):
         %0 = addf %a, %argb : f32
         linalg.yield %0 : f32
   } -> tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
+// CHECK-LABEL:   func @add_d_init(
+// CHECK-SAME:                     %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>>,
+// CHECK-SAME:                     %[[VAL_1:.*]]: f32) -> tensor<32xf32> {
+// CHECK:           %[[VAL_2:.*]] = constant 32 : index
+// CHECK:           %[[VAL_3:.*]] = constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_4:.*]] = constant 0 : index
+// CHECK:           %[[VAL_5:.*]] = constant 1 : index
+// CHECK:           %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf32>
+// CHECK:           %[[VAL_7:.*]] = memref.alloc() : memref<32xf32>
+// CHECK:           linalg.fill(%[[VAL_3]], %[[VAL_7]]) : f32, memref<32xf32>
+// CHECK:           scf.for %[[VAL_8:.*]] = %[[VAL_4]] to %[[VAL_2]] step %[[VAL_5]] {
+// CHECK:             %[[VAL_9:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_8]]] : memref<?xf32>
+// CHECK:             %[[VAL_10:.*]] = addf %[[VAL_9]], %[[VAL_1]] : f32
+// CHECK:             memref.store %[[VAL_10]], %[[VAL_7]]{{\[}}%[[VAL_8]]] : memref<32xf32>
+// CHECK:           }
+// CHECK:           %[[VAL_11:.*]] = memref.tensor_load %[[VAL_7]] : memref<32xf32>
+// CHECK:           return %[[VAL_11]] : tensor<32xf32>
+// CHECK:         }
+func @add_d_init(%arga: tensor<32xf32, #DV>, %argb: f32) -> tensor<32xf32> {
+  %u = linalg.init_tensor [32] : tensor<32xf32>
+  %0 = linalg.generic #trait1
+     ins(%arga: tensor<32xf32, #DV>)
+    outs(%u: tensor<32xf32>) {
+      ^bb(%a: f32, %x: f32):
+        %0 = addf %a, %argb : f32
+        linalg.yield %0 : f32
+  } -> tensor<32xf32>
+  return %0 : tensor<32xf32>
+}
+
 // CHECK-LABEL:   func @mul_d(
 // CHECK-SAME:                %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>>,
 // CHECK-SAME:                %[[VAL_1:.*]]: f32,
 // CHECK-SAME:                %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
 // CHECK:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK:           %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf32>