[AMD] Introduce specialized Allocation pass #7328
Conversation
This PR introduces AMD specific allocation pass and new attribute that controlls method scratch pad conversion method: padded or swizzled. For now OptimizeLDSUsage pass sets all convert_layout operations in padded mode.
|
General swizzling conversion consumes a lot more shared memory, which is a problem on mi30x and older architectures. Idea is to support both variants in AMD backend:
This PR adds only part related to allocation analysis. Conversion pattern implementation is in progress. |
| } | ||
|
|
||
| unsigned AMDAllocationAnalysisScratchSizeFn(Operation *op) { | ||
| if (op->hasAttr(AttrSharedMemPadded)) { |
There was a problem hiding this comment.
wouldn't that affect the codegen? I don't see any changes there?
There was a problem hiding this comment.
For AMD we still use old pattern, which is using padded memory, so this should be safe.
|
@ThomasRaoux hi
This particular PR do not affect codegen, it affects only allocation analysis. Analysis pessimistically allocates maximum memory from swizzled and padded memory: https://github.com/triton-lang/triton/blob/main/lib/Analysis/Allocation.cpp#L206 Codegen use old padded pattern for AMD: https://github.com/triton-lang/triton/blob/main/lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp#L287 |
I am not sure this should be in common code. It seems NVidia backend is mostly using swizzling for everything and gets rid of old pattern, that do not fit linear layout. +cc @antiagainst |
ah I see, so the control should ensure that we use the padded pass in the lowering though? Otherwise we just rely on implicit assumptions. |
|
Yes, in the nvidia backend we are never using padding. We always use swizzling or the stmatrix pass (which I'm working to integrate in the swizzling pass). I would suggest AMD also uses swizzling if possible given that it's more efficient in general and it will allow to generate swizzled layouts that can be lowered with special instructions like ldmatrix/stmatrix (not sure if AMD has special instructions like these). If this is not possible, AMD might want to end up moving the padding codegen to the AMD folder as this will not be needed for nvida. |
| auto scratchConfig = getScratchConfigForCvt(srcTy, dstTy); | ||
| elems = getNumScratchElements(scratchConfig.paddedRepShape); | ||
| } else { | ||
| // TODO use swizzling |
There was a problem hiding this comment.
assert that this path is not taken?
Right, this is what I want to implement for now. Keep padding related code in AMD backend, at least until we manage to lower memory consumption of swizzling pattern. |
|
|
||
| namespace mlir::triton::gpu { | ||
|
|
||
| void fillAllocationInfo(ModuleOp mod, ModuleAllocation &allocation); |
There was a problem hiding this comment.
Maybe name it as attachAllocationSizeAndOffsetAttr is clearer. Also can you add documentation to this function?
| #blocked1 = #ttg.blocked<{sizePerThread = [8, 4], threadsPerWarp = [8, 8], warpsPerCTA = [4, 1], order = [1, 0]}> | ||
| #blocked2 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [8, 8], warpsPerCTA = [4, 1], order = [1, 0]}> | ||
|
|
||
| // CHECK: ttg.shared = 36864 : i32 |
There was a problem hiding this comment.
Can you explain in comments how this number is computed so easier to understand/update later?
| #blocked1 = #ttg.blocked<{sizePerThread = [8, 4], threadsPerWarp = [8, 8], warpsPerCTA = [4, 1], order = [1, 0]}> | ||
| #blocked2 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [8, 8], warpsPerCTA = [4, 1], order = [1, 0]}> | ||
|
|
||
| // CHECK: ttg.shared = 131072 : i32 |
|
|
||
| constexpr char AttrSharedMemPadded[] = "amdgpu.shared_mem_padded"; |
There was a problem hiding this comment.
kPaddedScratchShmemAttrName[] = amdgpu.use_padded_scratch_shmem to be precise?
This PR introduces AMD specific allocation pass and new attribute that defines conversion method: padded or swizzled.
For now OptimizeLDSUsage pass sets all convert layout operations in padded mode.