[mlir][sparse] Implement BufferizableOpInterface for additional ops

[springerm]

Summary

The handling of unknown ops will be tightened in a subsequent change. All sparse_tensor ops should implement BufferizableOpInterface, otherwise, they are treated as "unknown" and additional buffer allocs/copies may be inserted around them.

Depends On: D142728

[Peiming]

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

190

The ToIndicesOp (as well as following ToXXX) returns a field of the sparse tensor, I am not sure whether it should be considered as an alias to the sparse tensor or not.

[springerm]

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

190

The result is a memref. getAliasingOpResult/getAliasingOpOperand model relationships between tensors. So this has to return {}.

But the question is: Should bufferizesToMemoryWrite return true? If someone writes into the result of to_indices, the sparse tensor is modified. This op is a boundary between tensors and memrefs, similar to bufferization.to_memref. It should probably bufferize to a memory write to account for all potential writes that could happen to the memref.

How are these sparse_tensor.to_... ops used? Are these ops that the user may write? Are there cases where you write into the returned memrefs?

[Peiming]

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

190

I think these ops are unlikely to be created by users, it is a relatively low-level operations inside sparse pipeline.

We do write into the memref returned, but only when initializing the sparse tensor (e.g., read from files, during concatenation/conversion, etc), but currently not after the sparse tensor are fully "loaded".

@aartbik (just to make sure you saw the revision when you are back and please correct me if I am wrong).

[aartbik]

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

190

In code generated by the sparse compiler, we indeed will never write to the memrefs returned by these operations directly (pointers/indices, for values we actually may when modifying values in place), but "user" code could of course write to the pointers and indices. Should we document that as undefined behavior and use "false"? What would the consequence be of using "true" conservatively?

[springerm]

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

190

Conservatively returning true would be the right thing to do. But it would introduce additional copies. E.g., this a copied from Integration/.../sparse_abs.mlir:

...
    %x = sparse_tensor.values %0 : tensor<?xf64, #SparseVector> to memref<?xf64>
    %y = sparse_tensor.values %1 : tensor<?xi32, #SparseVector> to memref<?xi32>
    %a = vector.transfer_read %x[%c0], %df: memref<?xf64>, vector<12xf64>
    %b = vector.transfer_read %y[%c0], %di: memref<?xi32>, vector<9xi32>
    %na = sparse_tensor.number_of_entries %0 : tensor<?xf64, #SparseVector>
    %nb = sparse_tensor.number_of_entries %1 : tensor<?xi32, #SparseVector>
...

There could be IR that writes into %x which would modify the buffer. Therefore, bufferizesToMemoryWrite = true would be correct.

But then there is a read of %0: sparse_tensor.number_of_entries %0. So the bufferization would make a copy of %0 before passing it into sparse_tensor.values, which is probably not desirable.

Note: sparse_tensor.number_of_entries reads just the number of elements and writing to sparse_tensor.values cannot modify the number of elements. That would be fine. But we don't have such a fine-grain level of read/write effects in the bufferization. Let's imagine that there's a linalg.generic ins(%0) after the sparse_tensor.values %0. Now we have an actual problem.

Should we document that as undefined behavior and use "false"?

That would be a good solution from a bufferization perspective.

[springerm]

add missing op interface

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

 
 using namespace mlir::bufferization;
 using namespace mlir::sparse_tensor;
 
 namespace mlir {
 namespace sparse_tensor {
 namespace {
 
+struct ConcatenateOpInterface
+    : public BufferizableOpInterface::ExternalModel<
+          ConcatenateOpInterface, sparse_tensor::ConcatenateOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+
+  bool isWritable(Operation *op, Value value,
+                  const AnalysisState &state) const {
+    return true;
+  }
+};
+
 struct ConvertOpInterface
     : public BufferizableOpInterface::ExternalModel<ConvertOpInterface,
                                                     sparse_tensor::ConvertOp> {
   bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
     // ConvertOps may allocate. (Unless they convert between two identical
     // types, then they fold away.)
     return true;
   }
struct InsertOpInterface
   BufferRelation bufferRelation(Operation *oo, OpResult opResult,
                                 const AnalysisState &state) const {
     // InsertOp returns the same object (realloc should not invalidate
     // aliases).
     return BufferRelation::Equivalent;
   }
 };
 
+struct NumberOfEntriesOpInterface
+    : public BufferizableOpInterface::ExternalModel<
+          NumberOfEntriesOpInterface, sparse_tensor::NumberOfEntriesOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+};
+
+struct ToIndicesBufferOpInterface
+    : public BufferizableOpInterface::ExternalModel<
+          ToIndicesBufferOpInterface, sparse_tensor::ToIndicesBufferOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // Potential writes into memory through the result of sparse_tensor.indices
+    // are not considered.
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+};
+
+struct ToIndicesOpInterface
+    : public BufferizableOpInterface::ExternalModel<
+          ToIndicesOpInterface, sparse_tensor::ToIndicesOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // Potential writes into memory through the result of sparse_tensor.indices
+    // are not considered.
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+};
+
+struct ToPointersOpInterface
+    : public BufferizableOpInterface::ExternalModel<
+          ToPointersOpInterface, sparse_tensor::ToPointersOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // Potential writes into memory through the result of sparse_tensor.pointers
+    // are not considered.
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+};
+
+struct ToValuesOpInterface
+    : public BufferizableOpInterface::ExternalModel<ToValuesOpInterface,
+                                                    sparse_tensor::ToValuesOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // Potential writes into memory through the result of sparse_tensor.values
+    // are not considered.
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+};
+
 } // 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);
-      });
+  registry.addExtension(+[](MLIRContext *ctx,
+                            sparse_tensor::SparseTensorDialect *dialect) {
+    sparse_tensor::ConcatenateOp::attachInterface<ConcatenateOpInterface>(*ctx);
+    sparse_tensor::NumberOfEntriesOp::attachInterface<
+        NumberOfEntriesOpInterface>(*ctx);
+    sparse_tensor::ToIndicesBufferOp::attachInterface<
+        ToIndicesBufferOpInterface>(*ctx);
+    sparse_tensor::ToIndicesOp::attachInterface<ToIndicesOpInterface>(*ctx);
+    sparse_tensor::ToPointersOp::attachInterface<ToPointersOpInterface>(*ctx);
+    sparse_tensor::ToValuesOp::attachInterface<ToValuesOpInterface>(*ctx);
 }