This is an archive of the discontinued LLVM Phabricator instance.

[mlir][sparse] Implement BufferizableOpInterface for additional ops
ClosedPublic

Authored by springerm on Jan 18 2023, 3:46 AM.

Details

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

Diff Detail

Event Timeline

springerm created this revision.Jan 18 2023, 3:46 AM
Herald added a project: Restricted Project. · View Herald TranscriptJan 18 2023, 3:46 AM
springerm requested review of this revision.Jan 18 2023, 3:46 AM
Peiming added inline comments.
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 added inline comments.Jan 19 2023, 9:56 AM
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 added inline comments.Jan 19 2023, 11:05 AM
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 added inline comments.Jan 23 2023, 5:28 PM
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 added inline comments.Jan 24 2023, 2:06 AM
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 edited the summary of this revision. (Show Details)Jan 27 2023, 8:13 AM
springerm updated this revision to Diff 492784.Jan 27 2023, 8:49 AM

add missing op interface

aartbik accepted this revision.Jan 27 2023, 8:50 AM
This revision is now accepted and ready to land.Jan 27 2023, 8:50 AM
This revision was landed with ongoing or failed builds.Jan 27 2023, 8:58 AM
This revision was automatically updated to reflect the committed changes.