[RFC] Memref index curry

MLIR
1

Something some others and I have come across is a need for an index curry for memref types. By this specifically I mean that I would like to be able to perform the first level of indexing to a load/store prior to the load/store operation

This is somewhat similar to the subview op, but not exactly the same. In particular subview requires knowing the remaining size/bounds/access pattern, whereas this is not necessarily always available in our use case.

A simple example of this op would be as follows:

%idx = ... : index
%ref = ... : memref<?x5xi32>
%sref = subindex %ref %idx : memref<?x5xi32> to memref<5xi32>
%loaded = load %sref[%c3]

A second, related use case is performing an offset on the first index (again without knowing the remaining bounds). This could be a different operation, but for ease I’ve also demonstrated it using “subindex” again.

%idx = ... : index
%ref = ... : memref<?xi32>
%sref = subindex %ref %idx : memref<?xi32> to memref<?xi32>
%loaded = load %sref[%c3]; // perform a load at 3 + %idx

I wonder if people have seen similar issues and generally if they have thoughts on this.

For our use case, we’ve implemented a simple version of the above “subindex” that does the corresponding operations, with relevant canonicalizations that attempts to remove the subindex when it sees a load/store/etc.

cc @ftynse @lorenzo_chelini @kumasento

1 Like
2

Yes we definitely want to push subview/subtensor and friends in that direction but we probably need to start with tonight’s ODM discussion on sparse type representation.

This will start pushing us in inspector-executor territory so thinking more generally in terms of n-D slices on non-dense data types will be useful. This quickly gets into the territory of computing a min/max over k-D slices at a dynamic position and into memref<memref> territory.

So far we have been able to emulate enough behavior with pairs of 1-D memref + op semantics but we are reaching the point where we want the next design step function.

Do you have particular data types in mind beyond sparse ?

3

Hi Nicolas,
I don’t think we have a specific data type in mind (@wsmoses please correct me if I am wrong). We would like to index a specific offset inside a memref without knowing the bounds. We could perhaps put some code together and create an RFC, what do you think?