# Design for direct load to LDS

## Short Introduction

### Goal
We try to implement a way to efficiently generate `global_load_lds` intrinsic to achieve direct loads from global memory to LDS memory, bypassing the need for registers.

### Target instruction: `global_load_lds`/`scratch_load_lds`
Note: for now we only consider `global_load_lds` instruction.
Although not very well explained in the AMDGPU programming manual, `global_load_lds` loads directly from global memory space on to LDS.
* The following operands are important factors:
* source global memory address and destination LDS address;
* bitwidth of the load, only a few variants are supported. On MI300X, 4, 8, 16 bit widths are supported.
* It is a subgroup instruction:
* meaning the LDS address needs to be coalesced across subgroup;
* global memory addresses within a subgroup do not need to be coalesced;
* In such a way, `global_load_lds` instruction works as a gather instruction.
* HIP example:

thread_specific_global_address = compute_thread_specific_load_address(threadIdx);
__builtin_amdgcn_global_load_lds(thread_specific_global_address,
                                 subgroup_specific_LDS_address,
                                 /*load width =*/4, 0, 0);

The above intrinsic instructs each thread (with subgroup ID = `subgroupIdx`, ranging `[0, 64)`) to load `4` bytes from `thread_specific_global_address`, to `subgroup_specific_LDS_address[subgroupIdx]`, resulting a consecutive `64 * 4 = 256`-byte memory chunk at `subgroup_specific_LDS_address`. Load width could be either 1, 2, 4 on MI300X, resulting 64, 128, 256 byte memory chunk being loaded.

### Using `global_load_lds` to load/gather large chunks for each thread

A subgroup of threads must coordinate to gather a consecutive chunk due to limitation. Here is an example:
In the example, thread 1 is unable to load a consecutive chunk itself, and it has to load strided values, and it has to orchastrate with all the threads in the subgroup.

Here is an HIP example of a 1-D `memcpy` within a workgroup:

// Loading `size` of bytes from `global` to `shared`.
__device__ void chunk_memcpy(__shared__ uint8_t* shared, uint8_t *global, size_t size) {
  const size_t subgroup_size = 64;
  const size_t load_width = 4;
  const size_t subgroup_copy_size = load_width * subgroup_size;

  auto lane_id = workgroup_thread_id % subgroup_size;
  auto subgroup_id = workgroup_thread_id / subgroup_size;

  auto subgroup_shared_addr = shared[(subgroup_id * subgroup_size) * load_width];
  auto subgroup_global_addr = global[workgroup_thread_id * load_width];
  auto copy_size = workgroup_thread_size * load_width;

  // Assuming size % total_copy_size == 0 for simplicity
  for (size_t i = 0; i < size; i += copy_size) {
 __builtin_amdgcn_global_load_lds(subgroup_global_addr[i], 
                                  subgroup_shared_addr[i],
                                  4, 0, 0);
  }
}

* Each thread gathers from a different address in global, but a subgroup shares a same destination address.

## Proposed Changes to existing codebase:

### New IREE_GPU attribute `iree_gpu.use_global_load_dma` for `linalg.copy`
* In `GPUPromoteMatmulOperandsPass`, by default a `linalg.copy` is emitted with `iree_gpu.derived_thread_config` attribute. Such as:

linalg.copy {#lowering_config = #iree_gpu.derived_thread_config}

This marks the copy will be distributed among threads;
* To mark the copy is a global memory direct load, instead we need to emit `iree_gpu.use_global_load_dma` attribute instead. Which will look like:

linalg.copy {#lowering_config = #iree_gpu.use_global_load_dma}

### New `iree_gpu.global_load_dma`, `amdgpu.gather_to_lds`, `rocdl.load_to_lds` opcodes
* We propose adding both `iree_gpu.global_load_dma` and `amdgpu.subgroup_gather_lds` ops.

#### `iree_gpu.global_load_dma`
* semantically, it copies elementwise a block of tensor via DMA. It represents a block-level DMA operation. The goal is to avoid tiled `linalg.copy` being further transformed when decisisons were made earlier.

* format:

Tensor = iree_gpu.global_load_dma ins(%src : tensor, %srcIndices) (outs: %dst : tensor) -> tensor {
  iree_gpu.yield %src[%srcIndices]
} 

* operands:
* `src`, `dst` and `result` are tensors.
* `$src[$srcIndices]` points to global memory space location for gathering, while `$dst[$dstIndices]` points to subgroup workgroup address space for storing.
* `src`, `dst` and `result` tensors must not have dynamic dims
* copy elementwise

* thread-distributed `iree_gpu.global_load_dma` will be mapped and lowered to `amdgpu.subgroup_gather_lds` in `LowerIREEGPUOpsPass`.

Here is how to emit `iree_gpu.global_load_dma`:
1. Usually, an emitted `linalg.copy{#lowering_config = iree_gpu.derived_thread_config}` in `GPUPromoteMatmulOperandsPass` will instead have `#iree_gpu.use_global_load_dma` attribute if using global loads is considered profitable.
2. A new pass right afer `GPUPromoteMatmulOperandsPass` called `GPULowerToGlobalLoadsPass` will decompose the tagged `linalg.copy` ops into a block of code which utilizes `iree_gpu.global_load_dma` op.

> [name=Quinn] This is why we need an op. We can't convert to `vector.transfer_read/write` ops because this operation is memory -> memory. transfer_read + write is memory -> register -> memory.

#### Example IR
Here is an IR example. To lower:

%9 = tensor.empty() : tensor<64x256xi8>
%10 = linalg.copy {lowering_config = #iree_gpu.use_global_load_dma} ins(%extracted_slice : tensor<64x256xi8>) outs(%9 : tensor<64x256xi8>) -> tensor<64x256xi8>

In `GPULowerToGlobalLoadsPass`, it is lowered to:

    %9 = tensor.empty() : tensor<64x256xi8>
    %10 = gpu.subgroup_id : index
    %11 = gpu.lane_id
    
    // extract global slice of the subgroup (probably don't need extract_slice) 
    %c16 = arith.constant 16 : index
    %c256 = arith.constant 256 : index
    %c4 = arith.constant 4 : index
    %12 = arith.muli %10, %c4 : index
    %extracted_slice_1 = tensor.extract_slice %extracted_slice[%12, 0] [16, 256] [256, 1] : tensor<64x256xi8> to tensor<16x256xi8>
    
    %c64 = arith.constant 64 : index
    %c4096 = arith.constant 4096 : index
    %13 = arith.muli %10, %c4096 : index
    %14 = arith.addi %13, %11 : index
    %c0_2 = arith.constant 0 : index
    %c64_3 = arith.constant 64 : index
    %c64_4 = arith.constant 64 : index
    
    // iterate over the global slice, and issue element-wise loads
    // onto workgroup tensor
    %15 = scf.for %arg3 = %c0_2 to %c64_3 step %c64_4 iter_args(%arg4 = %9) -> (tensor<64x256xi8>) {
      %25 = arith.muli %arg3, %c64 : index
      %26 = arith.addi %25, %14 : index
      %27:2 = affine.delinearize_index %26 into (16, 256) : index, index
      %28 = arith.addi %13, %25 : index
      %29:2 = affine.delinearize_index %28 into (64, 256) : index, index
      %30 = iree_gpu.global_load_dma %extracted_slice_1[%27#0, %27#1] -> %arg4[%29#0, %29#1] : tensor<16x256xi8>, tensor<64x256xi8> -> tensor<64x256xi8>
      scf.yield %30 : tensor<64x256xi8>
    }
    gpu.barrier

Comments:
* probably don't need to extract global tensor slice, and we can fold the indices into `iree_gpu.global_load_dma`.
* at the tensor level, using `iree_gpu.global_load_dma` must make sure that the target tensor will allocate a contiguous mem space, so probably no extract_slice is allowed.

#### `amdgpu.gather_to_lds` (add to upstream)
* This op semantically maps to the behavior of a rocdl `global_load_lds` intrinsic. It performs a subgroup-orchastrated gather from `src` (global, buffer, or other supported) memory to consececutive entries in LDS. It has subgroup-level semantics.

* format:

amdgpu.gather_to_lds $src `[` $srcIndices `]` `to` $dest `[` $destIndices `]` `:` $transferType `,` type($src) `,` type($dest) 

For example

amdgpu.gather_to_lds %global[%i0, %i1] to %lds[%j0, %j1] 
  : vector<2xf16>, memref<1024x1024xf16, #gpu.address_space<global>>, memref<256x256xf16, #gpu.address_space<workspace>>

* operands:
- `$src`: a global memref to copy from
- `$srcIndices`: the indices into the source memref. `$src[$srcIndices]` points to a gather source of `$transferType` (either scalar or vector), which is thread-specific.
- `$dest` - the shared memory memref to copy into
- `$destIndices` - the subgroup-level base index to copy into. Namely, `$dest[$destIndices]` defines the memref to the thread's subgroup's load target consecutive chunk in LDS memory.
- `$transferType` - the type that each subgroup will transfer. This determines the size of the copy. For example, `i32` or `vector<2xf16>` will both cause 4-byte copies.
- Size of the `$transferType` should be natively supported by `gloal_load_lds`, this is guaranteed when lowering `iree_gpu.global_load_dma`.
- element type of a vector must be the same as the memref.

### Lowering from `iree_gpu.global_load_dma` to `amdgpu.gather_to_lds`
* One `iree_gpu.global_load_dma` might be lowered to multiple `amdgpu.gather_to_lds` instructions, given that `gather_to_lds` has only single-word length transfer size as of now.

#### Example lowering IR:

%subview_4 = memref.subview %3[%69, %6] [1, 64] [1, 1] : memref<256x256xi8, #amdgpu.address_space<fat_raw_buffer>> to memref<1x64xi8, strided<[256, 1], offset: ?>, #amdgpu.address_space<fat_raw_buffer>>

%subview_5 = memref.subview %alloc_0[%arg4, 0] [1, 64] [1, 1] : memref<128x64xi8, #gpu.address_space<workgroup>> to memref<1x64xi8, strided<[64, 1], offset: ?>, #gpu.address_space<workgroup>>

iree_gpu.global_load_dma %subview_4, %subview_5 : memref<1x64xi8, strided<[256, 1], offset: ?>, #amdgpu.address_space<fat_raw_buffer>>, memref<1x64xi8, strided<[64, 1], offset: ?>, #gpu.address_space<workgroup>>

* In the above case, each thread is going to load `1x64xi8` which is 64 bytes, and it translates into 16 `global_load_lds` instructions. We could change tiling configurations to change how many `global_load_lds` instructions are being lowered.

#### Computing indices for `amdgpu.gather_to_lds`
##### Gathering source indices
* if the transfer type is scalar: Indices are incremented by 1 for each of the memref.
* if the transfer type is vector: Indices are incremented by the number of vector elements of the memref.
Note: this will probably not be efficient if multiple loads are needed.

##### Destination indices calculation (subgroup indices)
Discussion: There is no way at this late stage to get subgroup information. Either:
* when doing tiling, tile as subgroup info, and pass down subgroup info from `linalg.copy` to `iree_gpu.global_load_dma`, or
* create a new `amdgpu` op that can extract the information at runtime for the thread, however, we can only obtain the following values:
* workgroup size, subgroup size, thread id. It does not allow us to re-construct the destination's workgroup memref.

In sum: the `linalg.copy` should contain enough information to figure out the subgroup's memref. Here are some possible solutions:
* tile according to subgroup, so we just need to get the memref's head address.
* re-construct subgroup's memref offset by workgroup size, subgroup size.
* Preprequisite: 1. have a way to find those info; 2. tiled `linalg.copy`'s semantic is basically thread's moving, and the dest memref addresses are coalesced.

### New tiling configuration
Due to the limitation the subgroup instruction imposes, we must make sure that the tiling configuration in `lowering_config` attr must meet subgroup criteria. Consider the following factors:
* load byte width could be 1, 2 or 4 for MI300X, prefer 4 because of performance;
* Direct load instructions will always yield cosecutive chunks in shared memory. To simplify analysis, we ask that the load destination in shared memory to be consecutive.

This change would require to choose the tiling configuration according to the current configuration of workgroup sizes and subgroup sizes.

#### Loading multiple bytes per thread
* A thread is usually responsible loading multiple words.
* Naive tiling can only load up to a single word.

`GPUApplyTilingLevelPass` must decompose/update `linalg.copy` accordingly to accomondate mulitple-word memory loading.

##### Distributing data by threads
* each subgroup loads a contiguous region to allocated workgroup memory.
* For example:
* workgroup size 256, subgroup size 64, loading `<16x64xi32>`:
* slice up the destination tensor into `256 / 64 = 4`
* subgroup `Si` (out of 4) has slice `<4x64xi32>`, which translates into 1024 bytes, 256 words.
* each thread should emit 4 global load intrinsics
* thread `Ti` in subgroup `i` loads at address:
* `for j = 0 to 4: [j, Ti]`
* There are some restrictions
* workgroup memory size must be dividable because we are unlikely to allow padding

##### Index calcualtions
* `num_subgroups = workgroup_size / subgroup_size`
* getting global and workgroup tensor slice for `subgroup_idx`:
* slice along outermost dimension by `num_subgroups`, i.e. `[x / num_subgroups, y, z]`
* number of loads for a thread:
* `num_loads = tensor_slice_size / (load_width * subgroup_size)`
* A thread with `lane_idx` within a subgroup, will load at `indices` of global tensor slice:
* `linearized_idx = lane_idx + subgroup_load_size * i`, where `i = 0 to num_loads`
* `gather_indices = affine.delinearize_index(linearized_idx) into dims`, where `dims` are tensor slice dimension.
* Threads within a subgroup will store to workgroup tensor slice at `indices`:
* `linearized_idx = subgroup_load_size * i`, where `i = 0 to num_loads`
* `store_indices = affine.delinarize_index(linearized_idx) into dims`, where `dims` are tensor slice dimension.

##### Example IR:
* Before:

%empty = tensor.empty()
%copied = linalg.copy ins(%global) outs(%empty) ...

* After:

%empty = tensor.empty()
%tensor_out = scf.forall (%dim) -> {
    %global_slice = tensor.extract_slice %global[(see above section for details)] ...
    %shared_slice = tensor.extract_slice %empty[(see above section for details)] ...

    %index_0 = ...
    iree_gpu.global_load_dma %global_slice[%load_indices_0], %shared_slice[%store_indices_0]...
    %index_1 = ...
    iree_gpu.global_load_dma %global_slice[%load_indices_1], %shared_slice[%store_indices_1]...
    ...

    scf.yield %shared_slice
}

* Thoughts:
* bitcast slices to larger elements is possible, but we need to bitcast both global and shared slices.
* It is best to do it in `GPUApplyTilingLevelPass`.

###### Calculating indices
* global and shared memory slice for subgroups:
* given a tensor, slice it up to `workgroup_size / subgroup_size` slices.
*

## Other issues and considerations

### Gathers

The most important fusion to support will be gathers, meaning the representation of `iree_gpu.global_load_dma` ideally represents arbitrary gathers. The most straightforward way to do this is by modeling it the way a gather would typically be written with `linalg` and tap into the vectorizer.

### paddings?

IIUC this instruction doesn't support padding, so if we do need to do padding we will have to either do it on the load from LDS, or do it with supplementary (buffer)loads+stores. One option will be to do that on the conversion from `iree_gpu.global_load_dma` to `amdgpu`, but this only works when the inner most dimension of the matrix tile is *not* padded (typically K, so typically unpadded). In cases where multiple dimensions of the matrix tile are padded, it's unclear how coalescing the stores to shared memory will be possible outside of very specific shapes. This will be left as future work for the time being.

### Disabling fusion?

All fusions *must* happen after the load from shared, meaning to support fusions, we will need patterns to ensure these dma ops end up at the boundaries of the kernel. This is not the common case and can be left as future work.