Current memref type implementation contains explcit offset both as part of the type itself (memref<?xf32, strided<[1], offset: ?>>) and as part of its lowered data structure:
Most operations which change memref starting element (e.g. memref.subview) are keeping data pointer unchanged and instead manipulate offset field. Then data access ops (memref.load/store) combine data pointer with offset to get pointer to the actual memref contents. Such approach complicates memref lowering and was source of subtle bugs in the past without clear benefits. Also, offset removal was discussed multiple times over the years, but was never actually attempted, as far I know (e.g. [BUG][MLIR] expand-strided-metadata pass changes IR semantic · Issue #64334 · llvm/llvm-project · GitHub).
I’m proposing following changes:
Remove offset from lowered memref structure, any operations which need to change starting element will change data pointer instead.
Remove offset from StridedLayoutAttr and related infra. Memrefs which type only differs with offset (e.g. after 1D subview with step 1) will now be of the same type.
memref.reinterpret_cast and memref.extract_strided_metadata will keep offset argument/result, but semantics will change. Memref returned from reinterpret_cast with non-zero offset will have this offset applied directly to input memref pointer. Offset, returned from extract_strided_metadata will always be zero, but I will keep it for symmetry. Such semantics will allow existing reinterpret_cast-based memref.subview decomposition to work without changes. Ideally, we need a dedicated memref.ptr_add op for such decomposition so we can drop offsets arguments entirely, but it’s a more long term goal.
Also, it was discussed previously to change how we store allocated pointer in memref, to store offset relative to data instead of second pointer. But it’s completely orthogonal to the current proposal. For now, I’m not going to change it as part of current proposal and still store 2 pointers. We may change it in the future, but it will be a separate proposal, and this new offset will be just an implementation detail and not be part of StridedLayoutAttr type.
Unfortunately, I don’t have a good way to decompose it, so it will be 2 big patches:
Memref lowering change, removing offset from lowered structure and changing impl of reinterpret_cast and extract_strided_metadata. StridedLayoutAttr will still keep its offset at this point, but it will be effecively no-op during lowering.
In general, I am supportive of this direction. The offset field doesn’t pay for itself. I’d prefer this version:
but I agree with the motivation to first remove offset from the layout without further changes. Since it would be rather large breaking change for IR-based tests, I’d like to see some mechanism/instructions helping downstreams transition. If there is a way to layer this by having, e.g., the parser accept the offset but warn that it is ignored for O(weeks), it may help.
Actually I have serious concerns with this approach. Why should the offset be dropped? Apart from LLVM dialect we don’t have a way to represent pointer offset in MLIR. So dropping that would seem to have serious implications
+1, thanks for pushing this! Nobody got the activation energy to implement the change end-to-end. I think you have a reasonable plan, let me know if you need help!
(at minima I’ll try to provide some review time to ensure steady progress).
This sounds like an abuse of memrefs for a completely unrelated purpose. I know, we abuse of memrefs, too, and we have discussed pointer representation in the past. I don’t think it can be used as a reason to not progress.
There’s a ptr dialect, could we move your use there?
No… there are serialization formats (like SPIR-V modeled in SPIR-V dialect) cannot do pointer math. So they only way to handle this is through memrefs that have offsets. I dont see what is an abuse here. If you come from a pure LLVM background, it might seem that way, but from MLIR we can target not just LLVM where you cant just do pointer offset cause you dont have raw pointers. So, I dont see what the issue with having offset is. We have a full system that works well for LLVM and non-LLVM lowering that works well with offsets. I dont see any issue with having offsets in memrefs.
We’re not going to do this if just hand waving over the impact to SPIR-V and IREE.
Removing the offset from memref descriptor has no ambiguity for me. +1 to that.
I understand the appeal of going further to do the same to the high level type. But we need to have a worked out plan for what this means. If this is a design flaw, it is a five year old one. More diligence, please.
Ivan works on the IREE team, and I kind of hurriedly said “ok” when he asked me about this. I realize now that I was thinking about the descriptor but had not fully thought through the high level type. We’re going to have a meeting among our team members on this. Open to others joining. Let us know if interest, and we’ll get an appropriate meeting.
I indeed introduced this ~5y ago under the presumption that this will help encode and carry alignment information in the type but this use has not materialized to my original assumptions.
I have witnessed the offset be a cause of confusions and issues in various places for no clear benefit in the cases I’ve seen. The “Nobody got the activation energy to implement the change end-to-end” by @mehdi_amini is the most apt representation of the reality I’ve been privy to so far.
I have not seen the implications on the SPIR-V world myself, any benefit there from the representation has been a lucky stroke and not originally foreseen.
I am interested in seeing examples where the offset is necessary to update my priors.
Apart from LLVM dialect we don’t have a way to represent pointer offset in MLIR
I think there was a somewhat widely accepted proposal on a ptr dialect; maybe it is a prerequisite?
In the meantime, ExtractStridedMetadataOp/ExtractAlignedPointerAsIndexOp + addi + ReinterpretCastOp is a representation, I believe (not necessarily a good one).
Indeed. I’m neither for nor against doing this and have also been party to some of the issues. But it’s a burned in thing at this point. The way forward isn’t to handwave at a proposal that no one even knows if it has been actioned to a usable state.
It’s likely there’s a path forward here. We just need to analyze more closely vs having a big bang change without a plan.
Note that the offset is not really a part of the memref type. It’s part of the StridedLayoutAttr, which is but one possible layout, though likely the most frequently used one in a Linalg-based flow. Other layouts are theoretically possible (and an affine map is still expected to be the common denominator of those).
I don’t see what stops us from lowering any memref to a pair “base + offset”, regardless of it having offset: ? as part of the layout. Is there reliance on a static value of the offset somewhere?
If there were to be a meeting, an ODM could be a good venue. It’s underused and is at a time that works across most time zones.
I wouldn’t say that resolving this right this minute is a high priority for us – good to run down but has been this way for a long time. Perhaps some additional due diligence and use a Thursday slot next week or the week after.
It’s been a few years since I’ve looked at this deeply, and I’m refreshing some distant state.
I think from our perspective, abuse or not, there are two real use cases for which there are limited good ways to express this today:
SPIR-V needs to come from a form with fat pointers (index+offset) or have a completely reliable way to repatriate these.
IREE only uses base+offset (dynamic) at its kernel boundary, never just individual pointers for individual args. This is the result of ABI optimization being done on the kernels before the memref world ever gets them. Lowering these to a base+offset memref is the most natural thing today and, indeed, is what is done.
More hand-wavy, we’ve got another type that is “memref like” in that it maintains “analysis level” offsets that must be legalized in some way before lowering to a final target as, effectively a raw pointer that does not support arithmetic, and there is a belief/memory that this did unblock certain high level optimizations. Only including this because Nicholas asked if a use case for this kind of thing ever emerged. That’s the best I can recall.
I do recall from some work on this several years ago, that there was no use case for a static offset in the type (Alex’s question). And I also recall that having the type differentiate between a dynamic offsetable memref and a zero offset memref was a real pain that felt like something was missing to properly represent this.
I remember thinking about this in terms of there needing to be a more fundamental thing to represent a base+offset and that if that existed, the existing memref high level optimization infra needed to be sound with respect to producer offsets by default (as opposed to only considering it when the type indicated that it is an “offsettable memref”). I also recall doing an experiment to see just how unsound it was if just making memref “producer offset agnostic” and the answer was quite unsound. I walked away believing there was a better way to do this but that reality was stuck between missing a bit of representation and the memref high level optimizations needing more work to always be able to not destroy these relationships in the default (non offsettable) layout case.
Sorry, that is as far as my memory extends on the topic. But I think it corroborates the point about there possibly being missing pre-requisites that need to be met to unravel this case.
Edit: I recall thinking “what would it take to make the memref type always offsettable?” And concluding that a few ops needed to be more rigorously defined/created to always express offset and that the llvm descriptor side should just drop offset in favor of pointer math, with is perfectly fine at that level. Kind of a similar conclusion as everyone, but I think we need to create the missing representation bits and ensure it is sound. Descriptor offset removal could happen now (that will be disruptive in a different way because I’m sure that someone has some type punning structure interop that will fall apart, but that is outside of my viewpoint).
Good. It brings hope (to me at least) that the solution we find will be potentially good overall.
@Hardcode84 brings important enough issues that I feel we need some action, and has a reasonable (to my view) roadmap to remove it clean from upstream. Downstream is a different question, but we won’t know until we try.
Sorry, by “llvm descriptor”, do you mean the “LLVM Dialect usage” or the C++ code in LLVM (struct)? Trying to parse if you mean “memrefs get lowered into pointer arithmetic in LLVM dialect / IR” or “replace the memref type offset in the struct with some pointer arithmetic”.
If the latter, then how is that different than having an offset parameter?
(PS: assume ignorance, I’m not well versed in the particular issues you guys are having, just want to make sure we can continue using memref in a reasonable way for our cases).
Thinking about this a little bit more, I think we might be going after the wrong thing here. The real issue is the LLVM dialects memref descriptor. It contains sizes and strides, etc. and carries that information as an actual LLVM struct into LLVM IR. We are essentially relying on LLVMs SROA + optimization to fold it all away. While LLVM does a decent job of that, we are paying unnecessary compile time cost here.
What we really need to be doing is to linearize all memrefs before we lower to LLVM dialect. Correct me if I am wrong, but LLVM does not even handle the sizes and strides natively (but it can handle offsets natively through GEPs). So I would think it would be better to (eventually) get rid of the sizes and strides fields of the Memref descriptor and keep the offset. Wouldnt that go a long way in achieving what the stated goal of removing the offset is.
We already have all the tools needed to do the linearization (Indeed we do this in IREE on backends that dont support raw pointer types like SPIR-V backend, etc.), and are also using it as part of sub-byte load/store emulation (you really cannot do sub-byte load/store emulation without linearizing first). Having a dedicated linearization pass will allow us to
Materialize all the indexing logic and MLIR and fold it away. If we use affine.delinearize_index in combination with the affine.linearize_index being added a lot of the dimensionality related map will just fold away.
The same linearization is actually happening during many passes/conversion. I mentioned the sub-byte emulation path earlier, but also SPIR-V conversion path implicitly does this in the type converter. We can move all of that complexity out of those passes/conversions and make it explicit in MLIR before running any of those passes.
I think Alex is saying the same thing as what I am saying above. Any memref could be lowered to base + offset and then you have much less complexity in the Memref Descriptor. It is literally a base + offset. If we are trying to fix a problem here this seems more like a long term fix to me.
cc @krzysz00 (who is adding the affine.linearize and affine.delinearize ops and folders.
He works on my team on my projects and we’re not going to immolate ourselves but I am supportive of fixing this. I just don’t think for something like this that has been a sub-par design for five years, I can manufacture the urgency to leap without looking and a good plan.
Typo. Meant the memref descriptor data type throughout. Not sure what an “llvm descriptor” would be.
The problem I have with the layout having an offset is that this part is not used for anything meaningful as far as I know, and is actively contributing to making the IR larger because memref<???, strided<..., offset: ?>> is not cannot be used where memref<???, strided<...>> is expected and vice versa. This is a rather minor annoyance for me and the churn from the change is high, so I wouldn’t necessarily sponsor it nor would I oppose somebody doing this properly.
My use cases do not have an issue with having the offset in the descriptor, i.e. the structure memref lowers to when going to LLVM IR. But there is a bug referenced in the opening post.
I do have an issue with having two pointers (allocated and actual) in the descriptor. This may be confusing alias analysis, though noalias seems to be acceptable there as long as we don’t read through the allocated pointer.
Historically, the descriptor with explicit strides was proposed because the same format was used by (Py)Torch and seemed to work well there. I don’t have more recent data points that would invalidate this unless for fully static sizes.
Also historically, the descriptor in LLVM IR is a struct because we needed to keep and pass around dynamic stride values. There was no support for 1:N type conversion and no SRoA in MLIR, and not enough arithmetic/index canonicalization rules for that to be efficient. We didn’t want to reinvent what LLVM already does. Now we may have all of that and may want to reconsider. That being said, I won’t blindly accept the compile-time argument. MLIR is more expensive to operate, especially the dialect conversion part. If we do even more dialect conversion (extra 1:N type conversion to expand memref into set of values currently in the descriptor at the index level, prior to lowering to the LLVM dialect where we have only few canonicalizations + CSE + DCE), I’m not sure it offsets doing more passes on LLVM IR before we do SRoA.
There was some initial work in that direction with the ops that materialize descriptor data as values, but we didn’t push until the end.
I’ll vote for getting rid of the extra pointer in the memref descriptor. IIRC, it’s there because malloc() isn’t guaranteed to return sufficiently alligned memory? I figure that’s a bit weird reason to carry around an extra pointer - and it blocks the ability to annotate important properties like noalias at the backend level
Given the existence of things like SPIR-V fat pointers (and things like buffer resources on AMD cards, which ptr addrspace(7) is a hack around), I think Mahesh has a point - we could, before going to backend, lower all memrefs to 1-D memrefs, and explicitly surface dynamic dimensions in the calling convention before lowering. Strides could be handled there as well
Having been involved in this discussion for many years, it is pretty clear in my mind that the people closer to what this stuff lowers to and runs on consistently rate the memref descriptor and lowerings to it as the thing that needs the most attention. Specific reasons vary, but they are numerous (see @krzysz00 , @Hardcode84, SPIR-V discussion, etc).
Then there is usually an expressed opinion that linearization should be done before lowering to LLVM to eliminate all of that. A couple of years back, multiple groups kind of half way got what they wanted in that direction but it was incomplete. I have a clear memory of the work that @Hardcode84 refers to in this area.
And then there is the rationalization of the high level memref type so that we don’t have distinct types for dynamic-offset and static-offset (incl the default/implicit static offset-0) variants. Slight digression: It is an open question in my mind whether the problem is the offsets or if it is that the transformations are trying to be made sound on both representations when we should really just call the dynamic-offset one the only one and fill in the representation gaps to make that work.
I kind of have a suspicion that the solve here takes care of all of those facets and breaks down if just trying to address one of them. I’d particularly like to know what keeps the linearization options from just being the way. Is it “just engineering” hard or are we missing something?
I’m willing to fund fixing this but would prefer to fix it completely and end this epic once and for all. Curious what would get us there.
(edit: refined my terminology for dynamic-offset and static-offset memref types)
but I am supportive of fixing this. I just don’t think for something like this that has been a sub-par design for five years, I can manufacture the urgency to leap without looking and a good plan.