[AMD] Optimize address increments for buffer loads in loops #8464
Conversation
This PR transfers address computation from offsets in buffer loads to base pointers, which reuces amount of required computations and lowers register pressure.
binarman
force-pushed
the
buffer_load_base_opt
branch
from
82cef71 to
18bdacd
Compare
alefimov-amd
force-pushed
the
buffer_load_base_opt
branch
from
18bdacd to
226f8ef
Compare
|
|
||
| LogicalResult matchAndRewrite(scf::ForOp forOp, | ||
| PatternRewriter &rewriter) const override { | ||
| LDBG("Analyzing ForOp for for offset pointer optimization: " << forOp); |
| // baseIncrement is a scalar which will be added to base pointer after | ||
| // optimization offsetInitialized is a value of offset on first loop iteration | ||
| // incrementOp is an operation that advances offset tensor | ||
| struct LoadData { |
There was a problem hiding this comment.
"load is a target buffer load" should also include buffer_load_to_local.
You may want to manually format multi-line comment for baseIncrement.
"offsetInitialized" should be "offsetInitializer"
It's better to name this struct "LoadInfo". Just like what we have in the pipeliner.
There was a problem hiding this comment.
Current implementation covers all buffer ops, so I ended up with BufferOpInfo
| } | ||
| Value offsetInitializer = forOp.getInitArgs()[offsetOperandNo]; | ||
| LoadData data = {loadOp, advanceStep, Value(), offsetInitializer, | ||
| incrementOp}; |
There was a problem hiding this comment.
Why not call createScalarIncrements here to fill in baseIncrement? So after this analyzeLoad function, you have all info needed for this load.
There was a problem hiding this comment.
I want to fully split analysis and transformation. I am afraid of situation when optimization creates some operations, but later heuristic decides to skip.
There are no such heuristic at the moment, but I can not be sure we will not need it.
| // Gather buffer loads which could be optimized | ||
| SmallVector<LoadData> loads; | ||
| collectLoads<triton::amdgpu::BufferLoadOp>(loads, forOp); | ||
| collectLoads<triton::amdgpu::BufferLoadToLocalOp>(loads, forOp); |
There was a problem hiding this comment.
Can we create an interface for all these buffer ops, also including buffer_atomic_rmw and cas? So we don't have to go over the whole for loop again and again.
There was a problem hiding this comment.
Included this change in this PR and created separate here: #8600
This PR introduces a common interface for buffer address operands.
alefimov-amd
force-pushed
the
buffer_load_base_opt
branch
from
8fcf1d2 to
c2eacf4
Compare
3 participants
This PR transfers address computation from offsets in buffer loads to base pointers, which reuses amount of required computations and lowers register pressure.