Analyzing Partial Writes

MLIR
#1

I’m looking at doing some analysis to detect when a series of partial writes combines to be equal to a larger write. Is there any existing technology in MLIR i can leverage for this?

Example:

%0 = memref.alloc() : memref<1x3x512x256xbf16, 1>
%1 = memref.subview %0[0, 0, 0, 0] [1, 3, 256, 256] [1, 1, 1, 1] : memref<1x3x512x256xbf16, 1> to memref<1x3x256x256xbf16, #map0, 1>
%2 = memref.subview %0[0, 0, 256, 0] [1, 3, 256, 256] [1, 1, 1, 1] : memref<1x3x512x256xbf16, 1> to memref<1x3x256x256xbf16, #map1, 1>
// write to %1
// write to %2

In this example, I’m looking to be able to detect that the two subviews combine to be totally overlapping with the original tensor.

This analysis would need to detect both a sequence of writes as well as a subview that is created dynamically within a loop.

Thanks,
ian

#2

Not that I know of but it should be relatively straightforward to write after the recent refactor of FlatAffineConstraints that @matthias-springer did.

#3

In the general case, you need to define an PresburgerSet that represents the large memref, and another PresburgerSet that represents the union of small memrefs mapped back to the large memref address space. (You may also need to check that smaller memrefs don’t overlap by checking that the intersection of IntegerSets is empty.) Your equality analysis then boils down to demonstrating that the two integer sets are equal or, alternatively, that their symmetric difference is an empty set. PresburgerSet is a generalization of IntegerSet that allows for non-convexity.

If you need to also reason about operations in (affine) loops, you can obtain access relations using the MemRefAccess helper struct (only works on affine read/write operations though) and apply that to the iteration domain of the loop to obtain the set of values accessed by the loop. As before, it is necessary to remap that set back to the address space of the original memref.

There are no utility functions reasoning about subview with affine expressions AFAIK, but the address transformation of a subview operation with static shapes is an affine map (the one appearing as the layout map if the original layout is identity), so it shouldn’t be terribly complex to write the remapping.

1 Like
#4

Thanks, this looks very promising. Is there any examples of how to generate these types? its a bit of mental overload for me trying to understand and use.

#5

Tests are usually a good start - https://github.com/llvm/llvm-project/blob/main/mlir/unittests/Analysis/PresburgerSetTest.cpp

#6

Thanks. i’ve been reading through this code. I’m not seeing how to build the Integerset and/or PresburgerSet from an AffineMap. Is there an example of this? It sounds like maybe there’s one for AffineLoad/AffineStore? Though my search didn’t come up where to find this.

#7

Here is a simple example that builds something similar to what you want (ignore the more complex part that calls foldUpperBoundsIntoConstraintsSet): HoistPadding.cpp - llvm/llvm-project - Sourcegraph

It should be straightforward to adapt to build the “set of points in a subview”.

Once you have 2 such sets you can look at how to intersect them, this should also be quite mundane but I have not yet tried it myself as I only needed “projections” so far.

#8

You can try using FlatAffineRelation for this. There are functions available to convert AffineMap to FlatAffineRelation: llvm-project/AffineStructures.h at a55c4ec1cee7683d9095327d9d33e7137ec25292 · llvm/llvm-project · GitHub. This will give you a relation from iteration domain to memory locations. You can put constraints on the iteration domain to restrict size the domain. The range of this relation represents the memory accesses by the subview. You can extract the FlatAffineConstraints representing this range by using getRangeSet on the obtained relation. Doing this for each of the subview operation should give you sets corresponding to the accesses. From here, you can look at PresburgerSet tests to understand how to take union of these two sets.

1 Like
#9

Thanks everyone. This looks promising. I now need to find the time to play with this :slight_smile: