[AMD] Support TDM store on gfx1250 #8392
Conversation
borontion
requested review from
antiagainst,
peterbell10,
ptillet and
zhanglx13
as code owners
third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td
Outdated
Show resolved
Hide resolved
| auto [group0, group1] = LLVM::AMD::createTDMDescriptor( | ||
| rewriter, loc, getTypeConverter(), elementType, blockShape, tensorShape, | ||
| tensorStride, offset, srcPtr, dstPtr, op.getPred(), numWraps, | ||
| padInterval, padAmount); |
There was a problem hiding this comment.
why can't we do that when we create the descriptor instead of doing it on the fly?
There was a problem hiding this comment.
We need to know the shared memory destination and offset from the load/store op to create a complete tdm descriptor.
It is possible to create a partially filled tdm descriptor when lowering the create tensor descriptor op and then update it later for load/store, but this turns out still require some effort (specifically we need to extract and update global memory pointer, shared memory pointer and tensor shape
There was a problem hiding this comment.
the tensor shape and smem layout should be known at descriptor creation time. The problem is that the offsets need to be combined with the base global address?
I think separating the invariant part of the descriptor will be important to make sure we don't create code sequence accidentally loop dependent.
Also this solution will implicitly rely on LICM and CSE cleaning things up, I think this is a bit risky as in some cases we need to disable LICM to avoid register pressure
There was a problem hiding this comment.
Yeah it does depend on LICM to clean up. I can switch to do it in the make tensor descriptor if so. @antiagainst Do you have other thoughts?
There was a problem hiding this comment.
is there is no downsides I think it would be better. It would also reduce liveranges right? as I assume the descriptor is more packed than the set of values?
There was a problem hiding this comment.
Yeah what Thomas said makes sense to me. And expect it to provide better control and generate better code too.
| Type globalPtrTy = ptr_ty(ctx, 1); | ||
|
|
||
| Value globalAddrLow = group0[2]; | ||
| Value globalAddrHigh = b.and_(group0[3], b.i32_val(0x7FFFFFFF)); |
There was a problem hiding this comment.
As chatted internally, this is okay to get started for now. But we may want to define proper struct to avoid such i64 pack/unpack if LLVM is confused about it.
| return {group0, group1}; | ||
| } | ||
|
|
||
| void fillTDMDescriptor(RewriterBase &rewriter, Location loc, |
There was a problem hiding this comment.
not a blocking comment but at some point I would like to understand why we need to some much arithmetic for every load/store and how expensive it is.
5 participants
This PR adds support for TDM store on gfx1250, following #8333. Exposes
tdm.async_loadthrough Gluon. Groups common TDM utilities for load/store.