[mlir][sparse] attach bufferizableOpInterface to InsertOp

[Peiming]

Summary

To fix D136286

[aartbik]

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

118

typo: alias

[Peiming]

fix typo

[springerm]

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

100–103

false is the default implementation, so this method could be removed. Note: Any reallocation of sparse data structures "under the hood" is not visible to the bufferization (and it does not care about). The only important thing is that the "handle" to the buffer (in the dense case that's the "memref" buffer) does not magically change.

107

I'm wondering if this should return true.

Is the following IR valid?

%0 = bufferization.alloc_tensor
%1 = sparse.insert ... into %0
%2 = sparse.insert ... into %1

In that case, this should probably return true because there is a flow of information from the destination tensor to the result tensor.

But then there's also this op documentation:

Note that this operation is "impure" in the sense that even though
the result is modeled through an SSA value, the insertion is eventually
done "in place", and referencing the old SSA value is undefined behavior.
This operation is scheduled to be unified with the dense counterpart
`tensor.insert` that has pure SSA semantics.

What does that mean exactly? It's usually the bufferization's job to figure out implacability.

Are there certain restrictions wrt. copying a "not yet fully initialized" tensor? Also, is it considered an error if a sparse.insert would bufferize out-of-place. (We have BufferizableOpInterface::mustBufferizeInplace, which may or may not be useful here.)

130–133

This is also the default implementation in the op interface, so this method could be removed.

[springerm]

Sry, didn't see this revision in time, but I had a few comments.

[Peiming]

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

100–103

Okay

107

I think a new SSA value is required here simply because a realloc might happen (or not), which might invalidate the previous memory used for sparse tensors.

As you said, we should probably return true here, as to determine whether a reallocation is need a memory load is required for the current capacity (if that counts as a read).

I do not think we need mustBufferizeInplace here, because insert only creates realloc, which, as you said, does not matter to bufferization.

WDYT? @aartbik @bixia

[Peiming]

See D136510

[aartbik]

Thanks for having a look, Matthias!

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

107

Yes, for the short term. We want to migrate our sparse_tensor.insert to the tensor.insert in the longer term.
But for now we use this op in a very particular pattern:

%0= bufferization.alloc_tensor() : tensor<1024xf32, #SparseVector>
%1 = scf.for %i = %c0 to %c8 step %c1 iter_args(%vin = %6) -> tensor<1024xf32, #SparseVector> {

%ii = arith.muli %i, %c3 : index
%vout = sparse_tensor.insert %f1 into %vin[%ii] : tensor<1024xf32, #SparseVector>
scf.yield %vout : tensor<1024xf32, #SparseVector>

}
%2 = sparse_tensor.load %1 hasInserts : tensor<1024xf32, #SparseVector>

Right now basically only guarantee that the sparse vector is available at %2 but partially make sure we use SSA semantics for the construction. That is why we are avoiding the copy after %vout for now.

I wanted to talk to you about removing sparse_tensor.insert and go to the "pure" SSA version, so we should do that soon. In the meantime, this avoids a few complications for us.

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

bool isMemoryWrite(Operation *op, OpResult opResult,
     return false;
   }
 
   bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
     return true;
   }
 };
 
+struct InsertOpInterface
+    : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
+                                                    sparse_tensor::InsertOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    // Does bufferization handle realloc?
+    return false;
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // InsertOp writes to memory.
+    return true;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    // InsertOp returns an alias of its operand.
+    assert(op->getNumResults() == 1);
+    return op->getResults();
+  }
+
+  BufferRelation bufferRelation(Operation *oo, OpResult opResult,
+                                const AnalysisState &state) const {
+    // InsertOp returns the same object (realloc should not invalidate
+    // aliases).
+    return BufferRelation::Equivalent;
+  }
+
+  bool isWritable(Operation *op, Value value,
+                  const AnalysisState &state) const {
+    return true;
+  }
+};
+
 } // namespace
 } // namespace sparse_tensor
 } // namespace mlir
 
 void mlir::sparse_tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(
       +[](MLIRContext *ctx, sparse_tensor::SparseTensorDialect *dialect) {
         sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
         sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
         sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
+        sparse_tensor::InsertOp::attachInterface<InsertOpInterface>(*ctx);
       });
 }