[AMDGPU][LDS] Add in_bounds attribute to CoalescedGatherDMAOp for tensor.pad fusion

[lialan]

Summary

• Add in_bounds attribute to CoalescedGatherDMAOp to indicate which dimensions may have out-of-bounds reads
• Fuse tensor.pad into coalesced_gather_dma by tracing through to the original source tensor
• Enable coalesced DMA for unaligned matmuls that require padding

Motivation

When matmul dimensions don't align with tile sizes, tensor.pad is inserted to pad operands. Previously, this caused coalesced_gather_dma to fail because:

1. The padded tensor becomes a private memory allocation during bufferization
2. amdgpu.gather_to_lds requires source to be fat_raw_buffer or global memory

This PR fuses tensor.pad into the DMA operation, allowing it to read directly from buffer memory. AMD hardware's OOB behavior (returns 0 for out-of-bounds reads with boundsCheck=true) provides implicit padding.

Example

Before:

%extracted_slice_0 = tensor.extract_slice %6[%22, %19] [%17, %20] [1, 1] : tensor<65x121xf32> to tensor<?x?xf32>
%padded_3 = tensor.pad %extracted_slice_2 low[0, 0] high[%21, 0] {
      ^bb0(%arg4: index, %arg5: index):
        tensor.yield %cst : f32
} : tensor<?x64xf32> to tensor<4x64xf32>
%26 = tensor.empty() : tensor<4x64xf32>
%27 = linalg.copy {lowering_config = #iree_gpu.use_global_load_dma} ins(%padded_3 : tensor<4x64xf32>) outs(%26 : tensor<4x64xf32>) -> tensor<4x64xf32> 

We cannot handle such case because:

• When inferring address space of tensor.pad, it will always be assigned to private.
• global load DMA cannot handle private sources.

Notice that the semantics of the op is basically to extract a slice of an unaligned tensor to an aligned tensor. We can still do a DMA from the buffer pointer to utilize the address clamping, as long as the source is inferred to use fat raw pointer.

So the solution is to just fuse tensor.pad into coalesced_gather_dma when we are converting linalg.copy to iree_gpu.coalesced_gather_dma:

%extracted_slice_0 = tensor.extract_slice %6[%22, %19] [%17, %20] [1, 1] : tensor<65x121xf32> to tensor<?x?xf32>
%23 = tensor.empty() : tensor<32x4xf32>
%24 = scf.forall (%arg4, %arg5) = (0, 0) to (32, 4) step (32, 4) shared_outs(%arg6 = %23) -> (tensor<32x4xf32>) {
    %extracted_slice_7 = tensor.extract_slice %arg6[%arg4, %arg5] [32, 4] [1, 1] : tensor<32x4xf32> to tensor<32x4xf32>
    %30 = scf.forall (%arg7) in (64) shared_outs(%arg8 = %extracted_slice_7) -> (tensor<32x4xf32>) {
     ...
     %extracted_slice_9 = tensor.extract_slice %arg8[0, %arg7] [32, %33] [1, 1] : tensor<32x4xf32> to tensor<32x?xf32>
      ...
      scf.forall.in_parallel {
        iree_gpu.coalesced_gather_dma %extracted_slice_0 into %arg8 lane(%arg7) in_bounds [false, false] : tensor<?x?xf32>, tensor<32x4xf32>, index
      }
    } {mapping = [#iree_gpu.lane_id<0>]}
    scf.forall.in_parallel {
      tensor.parallel_insert_slice %30 into %arg6[%arg4, %arg5] [32, 4] [1, 1] : tensor<32x4xf32> into tensor<32x4xf32>
    }
} {mapping = [#gpu.warp<linear_dim_1>, #gpu.warp<linear_dim_0>]}

• The implicit fusing happens when we convert linalg.copy to iree_gpu.coalesced_gather_dma op (along with tiling).
• iree_gpu.coalesced_gather_dma now has attribute in_bounds which tells whether we are loading in bounds, so it now supports masking.
• make sure we source tensor is inferred to use buffer pointer.

[krzysz00]

I'd like to flag that we'll need to handle the slow path (address less than 4 bytes from the edge of the buffer)

I'll also note that we should check the semantics of gather to LDS within an if statement because it's possible this might work for global loads as well

[lialan]

I'll also note that we should check the semantics of gather to LDS within an if statement because it's possible this might work for global loads as well

@krzysz00 Did you mean we avoid using buffer read but stick with global loads? I am confused.

[krzysz00]

What I mean is that we should confirm that

x = ...; 
if (isInBounds(x)) {
  gather_to_lds global[x], lds[...] , ...
}

works in a way we can work with

If the answer is something like "yes, that causes the lanes that don't take the if to not write a value / write 0", then that'll be useful as a fallback for when we can't guarantee bufferability

[krzysz00]

Hold on this PR because I have correctness concerns

[qedawkins]

A few comments on top of what Krzysztof added.

[Yu-Zhewen]

Thanks! Some Nits.

[qedawkins]

I'm closing my eyes to the chipset leak. This should go in an AMD specific directory if it's going to be AMD specific.

[qedawkins]

This is the partial read off the end case. Does this work how we expect cc @krzysz00?