Page MenuHomePhabricator

[mlir] Support unranked memrefs, return and call ops in bare-ptr calling convention
ClosedPublic

Authored by dcaballe on Sep 15 2020, 2:01 PM.

Details

Summary

The bare-ptr calling convention lowers memref-typed values to LLVM
bare pointers at the boundaries of a function call. So far, we only
had support for ranked memref types and function arguments in this
convention. This patch adds support for unranked memref types and
'return' and 'call' ops. This mostly covers all the general use cases!

Some of these changes are also aimed at aligning the bare-ptr calling
convention code with the latest changes in the default calling convention
and reducing the amount of customization code needed. For example, the
'convertCallingConventionType' API is an attempt to abstract away the
conversion customization on types that are sensitive to the calling convention
and move those customizations from the conversion patterns to the type
converter. This is paving the way to isolate and even remove more calling
convention customizations in the future. For example, in follow-up
patches we should be able to unify 'barePtrFuncArgTypeConverter' and
'structFuncArgTypeConverter' into a single class. We could also create
independent type converters, one for each calling convention. Having
all these use cases working first would help us make sure that we preserve
the functionality along the improvements/refactoring.

Diff Detail

Event Timeline

dcaballe created this revision.Sep 15 2020, 2:01 PM
Herald added a project: Restricted Project. · View Herald TranscriptSep 15 2020, 2:01 PM
dcaballe requested review of this revision.Sep 15 2020, 2:01 PM

As a general observation, do you folks actually want the bare-pointer convention, or is it only there to attach noalias argument attribute? Maybe we should consider investing into expressing aliasing information in MLIR rather than keep the maintenance burden.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
376

I understand this code has been here before, but I am confused by the comment. The "shape can be computed statically" sounds like the code expects a static shape with sizes/strides/offset known at compile time. But getStridesAndOffset does _not_ guarantee that, all the results can be kDynamicStride...

378–381

isStrided does only the check

424

Nit: elide trivial braces

442

Could this rather create an LLVMIntegerType of the appropriate bitwidth? Then this function wouldn't need to be placed in the type converter.

1245

Nit: invert the condition + continue

1249–1251

Being able to pass an optional list of "skip these users" to replaceOp sounds useful indeed.

1258

Is there actually some guarantee that the type has static shape? Otherwise the call will just assert and fail on valid IR.

2200

Nit: the code around this uses unsigned

3453

Typo: one value

dcaballe updated this revision to Diff 293272.Sep 21 2020, 3:32 PM
dcaballe marked 4 inline comments as done.
  • Refactor code to a utility function.
  • Replace i64 type with index type.
  • Address Alex's feedback.

Thanks!

Thanks, Alex!

As a general observation, do you folks actually want the bare-pointer convention, or is it only there to attach noalias argument attribute? Maybe we should consider investing into expressing aliasing information in MLIR rather than keep the maintenance burden.

The noalias issue is one problem but I think we also need to model some kind of memref that is lowered to a bare pointer (e.g., if we want to pass a pointer to an external function). Creating specific dialects to cover each and every one of these special case doesn't scale well. I thought unranked memrefs had this goal but now I see that in the default calling convention they are lowered to another descriptor that also contains a field for the rank. I've also seen some memcopies being generated as part of the unranked memref lowering. I'm a bit confused. Do you know about the current use cases for unranked memrefs and if it's going to evolve somehow?

If we progressively have all the pieces needed, I'll be definitely happy to transition the bare pointer calling convention to whatever we see fit!

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
376

The "shape can be computed statically" sounds like the code expects a static shape with sizes/strides/offset known at compile time

Yes, that's actually the case (at least for know). I think the check was suggested in the previous review. I added the comment because it took me a while to understand the purpose. Would it make more sense to replace it with something like 'hasStaticShape'?

442

I can do that but we have other utilities doing similar transformation in this class: promoteOperands and promoteOneMemRefDescriptor. Would it make sense to keep all the type/descriptor conversion utilities here for consistency?

1258

Yes, that's what the check in convertMemRefToBarePtr is for. We only support static shapes in this convention.

The noalias issue is one problem but I think we also need to model some kind of memref that is lowered to a bare pointer (e.g., if we want to pass a pointer to an external function). Creating specific dialects to cover each and every one of these special case doesn't scale well.

What are those cases specifically? I don't see a problem with passing a memref descriptor or its unpacked form to an external function, other than being careful and using the proper ABI. Memref is not (just) a pointer by design, it also contains the rank/offset/sizes/strides information that we would just drop on the floor by treating memrefs as pointers. The closest analogy of a pointer is memref<? x type> or, equivalently memref<? x type, offset:0, strides:[1]>.

I thought unranked memrefs had this goal but now I see that in the default calling convention they are lowered to another descriptor that also contains a field for the rank.

Unranked memrefs were originally intended to move the rank-dispatch from function signatures into the function. Instead of having print_2d_memref_f32 and print_3d_memref_f32 with different signatures, we have just one print_memref_f32, which does switch(rank) internally.

How do you propose to support memrefs with dynamic rank ("unranked" is a terminology abuse IMO) if you don't store the rank?

I've also seen some memcopies being generated as part of the unranked memref lowering.

We need to allocate space to store shape/stride information of unknown-at-compile-time length. Inside a function, the allocation happens on stack to avoid expensive dynamic allocation, but if we need to return the descriptor, the allocation must outlive the function so it cannot be on stack anymore, hence a new allocation+copy (and copy+free on the caller's side). The memcpy we insert is an LLVM intrinsic, so it may just generate code for it rather than call the library.

Do you know about the current use cases for unranked memrefs and if it's going to evolve somehow?

Mostly API trick I described above. TF kernel generator (see recent ODS) also uses these combined with an external shape dialect to implement rank-polymorphic pointwise operations.

Nothing is likely to evolve unless there are concrete proposals or, at least, use cases that need to be addressed ;)

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
376

It depends on what are the actual requirements for the memref to be convertible to a bare pointer. hasStaticShape only checks that the shape is static, but does not guarantee that so are the offset and strides, or even that the memref has a strided form (it can have an arbitrary layout map). So if there some expectation of storage contiguity and zero offset, this code should check exactly for those. I'd work back from the indexing scheme.

442

Okay, I take the consistency argument here. In a longer term, we may consider factoring out all these utilities.

1258

Yeah, except that the check does not provide the guarantees required here.

Thanks, Alex! Sorry, I misunderstood what an unranked memref is modeling. Let me remove that from the patch since it's incorrect and only leave the support for call and return ops. We can add support for unranked memrefs when needed.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
376

Ok! I'm restricting this to memrefs with only static information, including shape, strides and offset. Note that there is no other specific assumption since the memref descriptor is recreated inside the function with all the static information.

dcaballe updated this revision to Diff 294219.Sep 24 2020, 7:40 PM
  • Removed support for unranked memrefs
  • Addressed feedback
ftynse accepted this revision.Tue, Sep 29, 4:46 AM

Please update the commit description to what it actually does now before landing.

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h
91

Typo: they -> the

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
392–393

Nit: prefer the early-return pattern for this

2762

Nit: the code actually extracts the aligned pointer (which is the right thing to do).

This revision is now accepted and ready to land.Tue, Sep 29, 4:46 AM
dcaballe marked 3 inline comments as done.Tue, Sep 29, 11:32 AM

Thanks, Alex. Will do.

This revision was landed with ongoing or failed builds.Tue, Sep 29, 12:09 PM
This revision was automatically updated to reflect the committed changes.