+1. I would expect user code to look like tensorType.getEncoding().isa<SparseTensorEncoding>() and not have a literal method isSparse on any interface.
1 Like
This looks strange to me. âClientsâ are interested in properties of tensors, not of attributes. Arguably, the tensor type interface implementation could go through an attribute interface internally (I started that way, but found the wrapper a bit too think to justify its existence), but the code above looks too raw to me.
This looks much better too me:
tensorType.isa<SparseTensor>()
tensorType.cast<SparseTensor>().getDimOrdering()
etc
I disagree. Clients are interested in the encoding of tensors. And encoding is an attribute. So accessing the attribute is not ârawâ but in fact expected. A large part of the design impetus for making this an attribute is to avoid the notion of tensorâs somehow being intrinsically able to be sparse, and replacing that with the encoding concept which generalizes better to an open ecosystem.
Iâd be ok with this if there is meaningful value to the SparseTensor class that is not captured in the SparseTensorEncoding. Doing this for pure syntactic convenience wrapping the SparseTensorEncoding seems like overkill, but I could be convinced.
Note that we can add a helper so that this syntactically works: tensorType.getEncoding<SparseTensorEncoding>() (it would be equivalent to dyn_cast_or_nullâing, like we do for value.getDefiningOp<FooOp>())
Arguably, there can be situations where it is convenient to query if a tensor is sparse without having to remember that sparsity is represented as encoding. For example, if one implements a code generation algorithm that works only for sparse tensors. Is there an established way to build such higher level abstractions that hide details of the representation? Or, is hiding details of the representation considered harmful?
This is true: but âsparseâ is not a property of the builtin tensor type today. We went with having the builtin tensor type only exposes an opaque encoding attribute, which put the âinterfaceâ as an interface attribute instead of a type attribute.
In order to have a âSparseTensorâ type interface that is implemented by the builtin tensor type, I believe we would need a âsparse encodingâ attribute interface that the builtin tensor type could target, acting as a âproxyâ between the âSparseTensorâ type interface and the attribute implementations.
That way you could get the API youâre looking for, without introducing any new coupling here.
You can have a SparseTensor class such that dyn_cast/isa will correctly cast to it only for tensor types with sparse encoding (the free function versions of isa/dyn_cast are highly customizable); you donât even need to involve TypeInterfaces or anything (just implement SparseTensor::classof correctly), and SparseTensor doesnât even have to inherit from Type at all. I donât think that is an âestablished wayâ, but it would solve the problem and might become an âestablished wayâ to deal with encodings 
What weâre trying to avoid is having TensorType or RankedTensorType have the union of methods for every possible encoding.
Note that we are also specifically trying to avoid introducing a new sparse tensor Type, which is a design we decided against from the outset.
I would be interested in looking at some passes written with the tensorType.getEncoding<SparseTensorEncoding>() API I suggested, and see if it is too bulky in practice / see what the pain points are.
Would the SparseTensor TypeInterface have any methods different from those on SparseTensorEncoding? If not, there doesnât seem much value in having it.
The only value I can see (beyond the slightly more concise API that Aart is looking for) is that weâd decouple the actual concrete type implementation from the interface (which isnât entirely free by the way), so that other non-builtin types can fit the interface and not necessarily have an attribute interface. It would also make all the sparse codegen infrastructure less coupled to the builtin tensor type.
(I donât claim to have an immediate use case for this)
Makes sense. Yeah, feels a bit light on use cases. One I can think of is that PyTorch has sparse tensors, but also its tensors are mutable, so having a SparseTensorTypeInterface that covers it would also imply that SparseTensorTypeInterface wouldnât allow its users to assume it is immutable.
Iâm still unclear on the syntactic benefit of manipulating SparseTensorTypeInterface interface object over using the SparseTensorEncoding interface object directly.
Would the advantage be that SparseTensorTypeInterface has the union of methods of RankedTensorType and SparseTensorEncoding? If so, that seems effectively identical to just introducing a new SparseTensorType that inherits from RankedTensorType. What is the purpose of having the encoding attribute if that is where we end up?
Yes, besides being concise, the direction I wanted to take with a type interface was indeed to abstract the sparse tensor interface from the âencodingâ implementation detail so that others would be free to implement it differently, yet have the same codegen path later on.
It seems I need to tread carefully to make sure we are all in agreement eventually. So perhaps the best next step is a small revision with just an attribute interface instead of the type interface.
2 Likes
I went back and forth a few times, but ended with a design that âfeels rightâ and hopefully is also acceptable to @mehdi_amini and @_sean_silva.
(1) The Tensor dialect introduces a new attribute sparse, which is basically just a wrapper around a dictionary.
#tensor.sparse< {dictionary} >
By itself, this attribute still has no intrinsic meaning, however, âŚ
(2) ⌠the new attribute implements a new attribute interface SparseTensorEncoding in the Tensor dialect, assigning meaning to the contents on the dictionary with the API described above (but no longer on tensor types). In addition, the interface defines a isValid() method which actually determines if the encoding makes sense for a given tensor.
So, now we are one step closer to defining a sparse tensor encoding, see an example below.
#CSR = #tensor.sparse< {
sparseDimLevelType = [ "dense", "compressed" ],
sparseDimOrdering = affine_map<(i,j) -> (i,j)>,
sparsePointerBitWidth = 64
sparseIndexBitWidth = 64
}>
(3) Now, remember the new encoding attribute in the TensorType in âcoreâ MLIR? By itself, this attribute also has no meaning. We can say things like
tensor<?x?xf64, "sparse">
tensor<?x?xf64, 12345>
tensor<?x?xf64, {a = b}>
and MLIR has no issue with that (other than asserting that these types are not identical).
But, we can also put the new tensor.sparse<> attribute as value!
tensor<?x?xf64, #CSR>
Now, for all most cases
tensorType.getEncoding().isa<SparseTensorEncoding>()
fails, except when we have the tensor.sparse<> attribute inside. Then we call isValid() on the corresponding interface (returned by cast<>) to make sure the encoding actually makes sense for the tensorType. After that the API of the attribute interface gives us âTACOâ style annotations back!
If needed, we can still introduce a type interface as well, but that is to be discussed another time.
I have a PR ready and will send this out shortly.
Overall sounds good.
two comments here:
-
It feels like we are missing some sort of verifier hook somewhere. Having potentially-invalid encoding attributes checked tediously for correctness by each pass seems wrong. It feels like we need an attribute interface VerifiableTensorEncoding with a âverifyâ method, and TensorTypeâs verifyConstructionInvariants could call
verifyif the encoding attribute implements VerifiableTensorEncoding. -
I donât see the value in having
SparseTensorEncodingbe an interface. It seems like having them just be methods of a concrete#tensor.sparseinstance is perfectly fine for now. (making it an interface requires a lot of assumptions about how this will be extended in the future, for example, that other tensor types will model sparsity as an encoding attribute â personally I would rather wait until we have concrete use cases to generalize this).
All of the code in tree that will manipulate directly the sparse attribute instead of the interface will limit the reuse later. That said I also suspect ideally the analyses and transformations would target a type interface instead of an attribute interface to provide the ability of someone using a sparse tensor type instead.
You may punt on this, but this would still required to find all the place in the code where youâll target the concrete instance and replace it with an interface later.
Yeah Iâd be in favor of requiring that every attribute used as encoding for a tensor implement such interface.
I guess what matters is the API for the interface, but is it the most efficient implementation for this attribute? For example sparseDimLevelType could be enum values instead of strings (even just stored as a bitfield). The ordering can be stored as a dense array of integer possibly. (This also opens up to custom printing/parsing)
It seems fairly easy to just grep for the concrete instance and replace with the interface name. If it is more involved than that, then by definition it means that clients are relying on concrete aspects of the #tensor.sparse attribute class that are not part of the interface. Thus prematurely putting those methods on the interface is not the design we want!
So either way it argues against prematurely introducing the interface. Either:
- The replacement is trivial, because the future interface has all the methods of the concrete attribute
- The replacement is nontrivial, which means that clients could not have been written against the interface in the first place (implying that we got the interface design wrong to begin with).
Note that I just revised the revision with a few changes for which attribute interfaces make more sense:
-
In âIRâ, we introduce two attribute interfaces:
VerifiableTensorEncoding: an API for verifying that encoding matches tensor type data
(used by parser and verifier to provide âsingleâ point verification)
SparseTensorEncoding: an API for a sparse tensor (nothing else), making it a first-class citizen -
in âTensor Dialectâ we introduce an concrete attribute
#tensor.sparse<dict>which implements both interfaces.
This is true, but only if we define #tensor.sparse with a strict interface: it shouldnât be exposing a dictionary and only has the same public method as what is exposed in the interface in the current revision (to be clear: Iâm fine if we do this instead of a separate interface from the concrete implementation).
+1, this is exactly what I imagined.
And just to whet your appetite on how close we are to properly typed sparse tensor code, consider how we currently represent sparse-matrix times vector (using lingalg specific annotations and glue, flags are required to define the overhead storage types):
!SparseTensor = type !llvm.ptr<i8>
#matvec = {
indexing_maps = [
affine_map<(i,j) -> (i,j)>, // A
affine_map<(i,j) -> (j)>, // b
affine_map<(i,j) -> (i)> // x (out)
],
sparse = [
[ "D", "S" ], // A
[ "D" ], // b
[ "D" ] // x
],
iterator_types = ["parallel", "reduction"],
doc = "X(i) += A(i,j) * B(j)"
}
...
func @kernel_matvec(%argA: !SparseTensor,
%argb: tensor<?xf64>,
%argx: tensor<?xf64>) -> tensor<?xf64> {
%arga = linalg.sparse_tensor %argA : !SparseTensor to tensor<?x?xf64>
%0 = linalg.generic #matvec
ins(%arga, %argb: tensor<?x?xf64>, tensor<?xf64>)
outs(%argx: tensor<?xf64>) {
^bb(%a: f64, %b: f64, %x: f64):
%0 = mulf %a, %b : f64
%1 = addf %x, %0 : f64
linalg.yield %1 : f64
} -> tensor<?xf64>
return %0 : tensor<?xf64>
}
...
%a = call @newSparseTensor(%fileName, %annotations, %u8, %u8, %f64)
: (!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
...
%0 = call @kernel_matvec(%a, %b, %x)
: (!SparseTensor, tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64>
With the type, this can now look as follows. This may not look like a very big difference, but it makes a huge difference in properly typing the operations!
#CSR = #tensor.sparse< {
sparseDimLevelType = [ "dense", "compressed" ],
sparseDimOrdering = affine_map<(i,j) -> (i,j)>,
sparsePointerBitWidth = 64,
sparseIndexBitWidth = 64
}>
#matvec = {
indexing_maps = [
affine_map<(i,j) -> (i,j)>, // A
affine_map<(i,j) -> (j)>, // b
affine_map<(i,j) -> (i)> // x (out)
],
iterator_types = ["parallel", "reduction"],
doc = "X(i) += A(i,j) * B(j)"
}
...
func @kernel_matvec(%arga: tensor<?x?xf64, #CSR>,
%argb: tensor<?xf64>,
%argx: tensor<?xf64>) -> tensor<?xf64> {
%0 = linalg.generic #matvec
ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?xf64>)
outs(%argx: tensor<?xf64>) {
^bb(%a: f64, %b: f64, %x: f64):
%0 = mulf %a, %b : f64
%1 = addf %x, %0 : f64
linalg.yield %1 : f64
} -> tensor<?xf64>
return %0 : tensor<?xf64>
}
...
%a = tensor.from_file %fileName : tensor<?x?xf64, #CSR>
...
%0 = call @kernel_matvec(%a, %b, %x)
: (tensor<?x?xf64, #CSR>, tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64>
2 Likes
Revision https://reviews.llvm.org/D101008 has been submitted!
After some internal discussion we decided to do the following:
(1) Keep the VerifiableTensorEncoding attribute interface in âIRâ, since the verification mechanism will be shared by all future encodings. Both the parser and tensor type verifier use this interface for âsingle pointâ verification.
(2) Define the SparseTensorEncoding attribute in the Tensor Dialect with a concrete interface for querying âTACOâ style sparse tensor formats.
The important thing is: we have a sparse tensor type in MLIR!
Next up: replace some of the clutter and glue in Linalg with the proper new sparse tensor type!
2 Likes
Thanks Aart for all the hard work and pushing this over the finish line! Great to see this new capability added 
1 Like