[LAYOUTS] [NFC] Make order accept a RankedTensorType #6007
Conversation
lezcano
requested review from
antiagainst,
zhanglx13 and
ptillet
as code owners
| // Order of the elements in the shared memory as defined at layout creation | ||
| // If this layout is associated with a MemDesc with a different shape | ||
| // it may return a different order than the actual order of the elements | ||
| SmallVector<unsigned> getDefaultOrder(SharedEncodingTrait layout); |
There was a problem hiding this comment.
Making this function accept an argument might be unclear. Cannot we just call SmallVector<unsigned> getOrder(MemDescType type);?
| @@ -134,8 +134,8 @@ sharedToLinearLayoutNoLeadingOffset(ArrayRef<int64_t> shape, | |||
| // Construct bases for the 2 most minor dimensions of the layout. These are | |||
| // the dims that get swizzled. | |||
| assert(shape.size() >= 2); | |||
| int colDim = shared.getOrder()[0]; | |||
| int rowDim = shared.getOrder()[1]; | |||
| int colDim = getDefaultOrder(shared)[0]; | |||
There was a problem hiding this comment.
For example, it's not clear to me why we do not call getOrder here
There was a problem hiding this comment.
As described in the OP, the idea is that getDefaultOrder should (in the future) just called in LinearLayoutConversions.cpp. It's basically just used in the definition of the LL.
For shared layouts, this may or may not be necessary, as we don't broadcast them over the size, so perhaps we could simply implement getOrder manually and call it a day.
I'm happy to do it that way, so that we just touch distributed for now.
lib/Dialect/TritonGPU/IR/Dialect.cpp
Outdated
| SmallVector<unsigned> getOrder(SharedEncodingTrait layout, | ||
| ArrayRef<int64_t> shape) { | ||
| return getDefaultOrder(layout); | ||
| if (auto swizzledLayout = | ||
| mlir::dyn_cast<SwizzledSharedEncodingAttr>(layout)) { | ||
| return llvm::to_vector(swizzledLayout.getOrder()); | ||
| } | ||
| if (auto sharedLayout = mlir::dyn_cast<NVMMASharedEncodingAttr>(layout)) { | ||
| return sharedLayout.getOrder(); | ||
| } | ||
| llvm::report_fatal_error("Unimplemented usage of getOrder for MemDescType"); | ||
| return {}; |
There was a problem hiding this comment.
I went this awkward way because I did not want to add a getOrder() API to SharedEncodingTrait, as I don't think it'll be useful in the future.
This is in preparation for moving the order to be implemented generically. We expose `getDefault.*Order` functions that implement hand-written order. We expect these functions to be used just in the LinearLayout creation. Eventually we'll inline them in that file and remove them completely.
This is in preparation for moving the order to be implemented
generically.
We expose
getDefault.*Orderfunctions that implement hand-writtenorder. We expect these functions to be used just in the LinearLayout
creation. Eventually we'll inline them in that file and remove them
completely.