// RUN: mlir-opt --fold-memref-alias-ops --split-input-file %s | FileCheck %s

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_subview_folding
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_subview_folding(%offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[ARG0]], %[[ARG1]]], %[[LOCAL]][%[[C0]], %[[C0]]]
  // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>
  %mem = memref.alloc() : memref<64x128xf16>
  %subview = memref.subview %mem[0, 0][32, 64][1, 1] : memref<64x128xf16> to memref<32x64xf16, strided<[128, 1]>>
  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %subview[%offset_i, %offset_j], %alloc[%c0, %c0]
    : vector<8xf16>, memref<32x64xf16, strided<[128, 1]>>, memref<64x64xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (s0 + 32)>
// CHECK: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 64)>

// CHECK: func @subview_folding_offset
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @subview_folding_offset(%offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: %[[IDX0:.*]] = affine.apply #[[MAP]]()[%[[ARG0]]]
  // CHECK: %[[IDX1:.*]] = affine.apply #[[MAP1]]()[%[[ARG1]]]
  // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[IDX0]], %[[IDX1]]], %[[LOCAL]][%[[C0]], %[[C0]]]
  // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>
  %mem = memref.alloc() : memref<64x128xf16>
  %subview = memref.subview %mem[32, 64][32, 64][1, 1] : memref<64x128xf16> to memref<32x64xf16, strided<[128, 1], offset: 4160>>
  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %subview[%offset_i, %offset_j], %alloc[%c0, %c0]
    : vector<8xf16>, memref<32x64xf16, strided<[128, 1], offset: 4160>>, memref<64x64xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_expand_shape
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_expand_shape(%offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[MEM:.*]] = memref.alloc() : memref<8192xf16>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: %[[IDXM:.*]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]]] by (64, 128) : index
  // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[IDXM]]], %[[LOCAL]][%[[C0]]]
  // CHECK-SAME: vector<8xf16>, memref<8192xf16>, memref<4096xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>
  %mem = memref.alloc() : memref<8192xf16>
  %expand_mem = memref.expand_shape %mem [[0, 1]] output_shape [64, 128] : memref<8192xf16> into memref<64x128xf16>
  %expand_alloc = memref.expand_shape %alloc [[0, 1]] output_shape [64, 64] : memref<4096xf16, #gpu_lds_addrspace> into memref<64x64xf16, #gpu_lds_addrspace>
  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %expand_mem[%offset_i, %offset_j], %expand_alloc[%c0, %c0]
    : vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_collapse_shape
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_collapse_shape(%offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>
  // CHECK: %[[INDICES_MEM:.*]]:2 = affine.delinearize_index %[[ARG0]] into (64, 128) : index, index
  // CHECK: %[[INDICES_LDS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (64, 64) : index, index
  // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[INDICES_MEM]]#0, %[[INDICES_MEM]]#1], %[[LOCAL]][%[[INDICES_LDS]]#0, %[[INDICES_LDS]]#1]
  // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>
  %collapse_alloc = memref.collapse_shape %alloc [[0, 1]] : memref<64x64xf16, #gpu_lds_addrspace> into memref<4096xf16, #gpu_lds_addrspace>
  %mem = memref.alloc() : memref<64x128xf16>
  %collapse_mem = memref.collapse_shape %mem [[0, 1]] : memref<64x128xf16> into memref<8192xf16>
  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %collapse_mem[%offset_i], %collapse_alloc[%offset_j]
    : vector<8xf16>, memref<8192xf16>, memref<4096xf16, #gpu_lds_addrspace>
  func.return
}


// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>


// CHECK: func @test_expand_shape_src_raw_buffer
// CHECK-SAME: %[[ARG0:.*]]: memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index
func.func @test_expand_shape_src_raw_buffer(%mem : memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: %[[IDXM:.*]] = affine.linearize_index [%[[ARG1]], %[[ARG2]]] by (64, 128) : index
  // CHECK: amdgpu.gather_to_lds %[[ARG0]][%[[IDXM]]], %[[LOCAL]][%[[C0]]]
  // CHECK-SAME: vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>
  %expand_mem = memref.expand_shape %mem [[0, 1]] output_shape [64, 128] : memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>> into memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>

  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %expand_mem[%offset_i, %offset_j], %alloc[%c0]
    : vector<8xf16>, memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_src_raw_buffer_collapse_shape
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>, %[[IDX:.*]]: index
func.func @test_src_raw_buffer_collapse_shape(%mem : memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>, %offset: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[IDX]] into (128) : index, index
  // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[INDICES]]#0, %[[INDICES]]#1], %[[LOCAL]][%[[C0]]]
  // CHECK-SAME: vector<8xf16>, memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>
  %collapse_mem = memref.collapse_shape %mem [[0, 1]]
    : memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>
    into memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>
  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %collapse_mem[%offset], %alloc[%c0]
    : vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>,
      memref<4096xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_expand_shape_dst_only
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_expand_shape_dst_only(%offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[MEM:.*]] = memref.alloc() : memref<8192xf16>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: %[[IDX_LDS:.*]] = affine.linearize_index disjoint [%[[ARG1]], %[[C0]]] by (64, 64) : index
  // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[ARG0]]], %[[LOCAL]][%[[IDX_LDS]]]
  // CHECK-SAME: vector<8xf16>, memref<8192xf16>, memref<4096xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>
  %mem = memref.alloc() : memref<8192xf16>
  %expand_alloc = memref.expand_shape %alloc [[0, 1]] output_shape [64, 64] : memref<4096xf16, #gpu_lds_addrspace> into memref<64x64xf16, #gpu_lds_addrspace>

  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds %mem[%offset_i], %expand_alloc[%offset_j, %c0]
    : vector<8xf16>, memref<8192xf16>, memref<64x64xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_nop
// CHECK-SAME: %[[ARG0:.*]]: memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index
func.func @test_nop(%mem : memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, #gpu.address_space<workgroup>>
  // CHECK: amdgpu.gather_to_lds %[[ARG0]][%[[ARG1]]], %[[LOCAL]][%[[ARG2]]]
  // CHECK-SAME: vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>
  amdgpu.gather_to_lds %mem[%offset_i], %alloc[%offset_j]
    : vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>

// CHECK: func @test_async_flag_preserved
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_async_flag_preserved(%offset_i: index, %offset_j: index) {
  // CHECK-DAG: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, #gpu.address_space<workgroup>>
  // CHECK-DAG: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK: amdgpu.gather_to_lds async %[[MEM]][%[[ARG0]], %[[ARG1]]], %[[LOCAL]][%[[C0]], %[[C0]]]
  // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu.address_space<workgroup>>

  %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>
  %mem = memref.alloc() : memref<64x128xf16>
  %subview = memref.subview %mem[0, 0][32, 64][1, 1] : memref<64x128xf16> to memref<32x64xf16, strided<[128, 1]>>
  %c0 = arith.constant 0 : index
  amdgpu.gather_to_lds async %subview[%offset_i, %offset_j], %alloc[%c0, %c0]
    : vector<8xf16>, memref<32x64xf16, strided<[128, 1]>>, memref<64x64xf16, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_transpose_load_subview
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_transpose_load_subview(%offset_i: index, %offset_j: index) -> vector<4xf16> {
  // CHECK-DAG: %[[ALLOC:.*]] = memref.alloc() : memref<64x128xf16, #gpu.address_space<workgroup>>
  // CHECK: amdgpu.transpose_load %[[ALLOC]][%[[ARG0]], %[[ARG1]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<workgroup>> -> vector<4xf16>

  %alloc = memref.alloc() : memref<64x128xf16, #gpu_wg>
  %subview = memref.subview %alloc[0, 0][32, 64][1, 1]
    : memref<64x128xf16, #gpu_wg> to memref<32x64xf16, strided<[128, 1]>, #gpu_wg>
  %result = amdgpu.transpose_load %subview[%offset_i, %offset_j]
    : memref<32x64xf16, strided<[128, 1]>, #gpu_wg> -> vector<4xf16>
  return %result : vector<4xf16>
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (s0 + 32)>
// CHECK: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 64)>

// CHECK: func @test_transpose_load_subview_offset
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_transpose_load_subview_offset(%offset_i: index, %offset_j: index) -> vector<4xf16> {
  // CHECK-DAG: %[[ALLOC:.*]] = memref.alloc() : memref<64x128xf16, #gpu.address_space<workgroup>>
  // CHECK: %[[IDX0:.*]] = affine.apply #[[MAP]]()[%[[ARG0]]]
  // CHECK: %[[IDX1:.*]] = affine.apply #[[MAP1]]()[%[[ARG1]]]
  // CHECK: amdgpu.transpose_load %[[ALLOC]][%[[IDX0]], %[[IDX1]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<workgroup>> -> vector<4xf16>

  %alloc = memref.alloc() : memref<64x128xf16, #gpu_wg>
  %subview = memref.subview %alloc[32, 64][32, 64][1, 1]
    : memref<64x128xf16, #gpu_wg>
    to memref<32x64xf16, strided<[128, 1], offset: 4160>, #gpu_wg>
  %result = amdgpu.transpose_load %subview[%offset_i, %offset_j]
    : memref<32x64xf16, strided<[128, 1], offset: 4160>, #gpu_wg> -> vector<4xf16>
  return %result : vector<4xf16>
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_transpose_load_expand_shape
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_transpose_load_expand_shape(%offset_i: index, %offset_j: index) -> vector<4xf16> {
  // CHECK-DAG: %[[ALLOC:.*]] = memref.alloc() : memref<4096xf16, #gpu.address_space<workgroup>>
  // CHECK: %[[IDX:.*]] = affine.linearize_index disjoint [%[[ARG0]], %[[ARG1]]] by (32, 128) : index
  // CHECK: amdgpu.transpose_load %[[ALLOC]][%[[IDX]]]
  // CHECK-SAME: memref<4096xf16, #gpu.address_space<workgroup>> -> vector<4xf16>

  %alloc = memref.alloc() : memref<4096xf16, #gpu_wg>
  %expand = memref.expand_shape %alloc [[0, 1]] output_shape [32, 128]
    : memref<4096xf16, #gpu_wg> into memref<32x128xf16, #gpu_wg>
  %result = amdgpu.transpose_load %expand[%offset_i, %offset_j]
    : memref<32x128xf16, #gpu_wg> -> vector<4xf16>
  return %result : vector<4xf16>
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_transpose_load_collapse_shape
// CHECK-SAME: %[[ARG0:.*]]: index
func.func @test_transpose_load_collapse_shape(%offset_i: index) -> vector<4xf16> {
  // CHECK-DAG: %[[ALLOC:.*]] = memref.alloc() : memref<32x128xf16, #gpu.address_space<workgroup>>
  // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[ARG0]] into (32, 128) : index, index
  // CHECK: amdgpu.transpose_load %[[ALLOC]][%[[INDICES]]#0, %[[INDICES]]#1]
  // CHECK-SAME: memref<32x128xf16, #gpu.address_space<workgroup>> -> vector<4xf16>

  %alloc = memref.alloc() : memref<32x128xf16, #gpu_wg>
  %collapse = memref.collapse_shape %alloc [[0, 1]]
    : memref<32x128xf16, #gpu_wg> into memref<4096xf16, #gpu_wg>
  %result = amdgpu.transpose_load %collapse[%offset_i]
    : memref<4096xf16, #gpu_wg> -> vector<4xf16>
  return %result : vector<4xf16>
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_transpose_load_nop
// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
func.func @test_transpose_load_nop(%offset_i: index, %offset_j: index) -> vector<4xf16> {
  // CHECK-DAG: %[[ALLOC:.*]] = memref.alloc() : memref<32x128xf16, #gpu.address_space<workgroup>>
  // CHECK: amdgpu.transpose_load %[[ALLOC]][%[[ARG0]], %[[ARG1]]]
  // CHECK-SAME: memref<32x128xf16, #gpu.address_space<workgroup>> -> vector<4xf16>
  // CHECK-NOT: subview
  // CHECK-NOT: expand_shape

  %alloc = memref.alloc() : memref<32x128xf16, #gpu_wg>
  %result = amdgpu.transpose_load %alloc[%offset_i, %offset_j]
    : memref<32x128xf16, #gpu_wg> -> vector<4xf16>
  return %result : vector<4xf16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_make_dma_base_subview
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index
func.func @test_make_dma_base_subview(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<64x64xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index) -> !amdgpu.tdm_base<f16> {
  // CHECK: %[[BASE:.*]] = amdgpu.make_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[LDS_I]], %[[LDS_J]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<64x64xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<f16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_base<f16>

  %subview = memref.subview %mem[0, 0][32, 64][1, 1] : memref<64x128xf16, #gpu_global_addrspace> to memref<32x64xf16, strided<[128, 1]>, #gpu_global_addrspace>
  %base = amdgpu.make_dma_base %subview[%global_i, %global_j], %lds[%lds_i, %lds_j]
    : memref<32x64xf16, strided<[128, 1]>, #gpu_global_addrspace>, memref<64x64xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_base<f16>
  func.return %base : !amdgpu.tdm_base<f16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_make_dma_base_expand_shape
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<4096xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index
func.func @test_make_dma_base_expand_shape(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<4096xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index) -> !amdgpu.tdm_base<f16> {
  // CHECK: %[[IDX:.*]] = affine.linearize_index disjoint [%[[LDS_I]], %[[LDS_J]]] by (64, 64) : index
  // CHECK: %[[BASE:.*]] = amdgpu.make_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[IDX]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<4096xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<f16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_base<f16>

  %expand_lds = memref.expand_shape %lds [[0, 1]] output_shape [64, 64] : memref<4096xf16, #gpu_lds_addrspace> into memref<64x64xf16, #gpu_lds_addrspace>
  %base = amdgpu.make_dma_base %mem[%global_i, %global_j], %expand_lds[%lds_i, %lds_j]
    : memref<64x128xf16, #gpu_global_addrspace>, memref<64x64xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_base<f16>
  func.return %base : !amdgpu.tdm_base<f16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_make_gather_dma_base_subview
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index
func.func @test_make_gather_dma_base_subview(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<64x64xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index) -> !amdgpu.tdm_gather_base<f16, i16> {
  // CHECK: %[[BASE:.*]] = amdgpu.make_gather_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[LDS_I]], %[[LDS_J]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<64x64xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_gather_base<f16, i16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_gather_base<f16, i16>

  %subview = memref.subview %mem[0, 0][32, 64][1, 1] : memref<64x128xf16, #gpu_global_addrspace> to memref<32x64xf16, strided<[128, 1]>, #gpu_global_addrspace>
  %base = amdgpu.make_gather_dma_base %subview[%global_i, %global_j], %lds[%lds_i, %lds_j]
    : memref<32x64xf16, strided<[128, 1]>, #gpu_global_addrspace>, memref<64x64xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_gather_base<f16, i16>
  func.return %base : !amdgpu.tdm_gather_base<f16, i16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_make_gather_dma_base_collapse_shape
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_IDX:.*]]: index
func.func @test_make_gather_dma_base_collapse_shape(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<64x64xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_idx: index) -> !amdgpu.tdm_gather_base<f16, i16> {
  // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[LDS_IDX]] into (64, 64) : index, index
  // CHECK: %[[BASE:.*]] = amdgpu.make_gather_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[INDICES]]#0, %[[INDICES]]#1]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<64x64xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_gather_base<f16, i16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_gather_base<f16, i16>

  %collapse_lds = memref.collapse_shape %lds [[0, 1]] : memref<64x64xf16, #gpu_lds_addrspace> into memref<4096xf16, #gpu_lds_addrspace>
  %base = amdgpu.make_gather_dma_base %mem[%global_i, %global_j], %collapse_lds[%lds_idx]
    : memref<64x128xf16, #gpu_global_addrspace>, memref<4096xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_gather_base<f16, i16>
  func.return %base : !amdgpu.tdm_gather_base<f16, i16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: #[[BOTH_MAP:.*]] = affine_map<()[s0] -> (s0 + 32)>
// CHECK: #[[BOTH_MAP1:.*]] = affine_map<()[s0] -> (s0 + 64)>

// CHECK: func @test_make_dma_base_both_fold
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<4096xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index
func.func @test_make_dma_base_both_fold(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<4096xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index) -> !amdgpu.tdm_base<f16> {
  // CHECK: %[[GI:.*]] = affine.apply #[[BOTH_MAP]]()[%[[GLOBAL_I]]]
  // CHECK: %[[GJ:.*]] = affine.apply #[[BOTH_MAP1]]()[%[[GLOBAL_J]]]
  // CHECK: %[[IDX:.*]] = affine.linearize_index disjoint [%[[LDS_I]], %[[LDS_J]]] by (64, 64) : index
  // CHECK: %[[BASE:.*]] = amdgpu.make_dma_base %[[MEM]][%[[GI]], %[[GJ]]], %[[LDS]][%[[IDX]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<4096xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<f16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_base<f16>

  %subview = memref.subview %mem[32, 64][32, 64][1, 1] : memref<64x128xf16, #gpu_global_addrspace> to memref<32x64xf16, strided<[128, 1], offset: 4160>, #gpu_global_addrspace>
  %expand_lds = memref.expand_shape %lds [[0, 1]] output_shape [64, 64] : memref<4096xf16, #gpu_lds_addrspace> into memref<64x64xf16, #gpu_lds_addrspace>
  %base = amdgpu.make_dma_base %subview[%global_i, %global_j], %expand_lds[%lds_i, %lds_j]
    : memref<32x64xf16, strided<[128, 1], offset: 4160>, #gpu_global_addrspace>, memref<64x64xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_base<f16>
  func.return %base : !amdgpu.tdm_base<f16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_make_dma_base_nop
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index
func.func @test_make_dma_base_nop(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<64x64xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index) -> !amdgpu.tdm_base<f16> {
  // CHECK: %[[BASE:.*]] = amdgpu.make_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[LDS_I]], %[[LDS_J]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<64x64xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<f16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_base<f16>
  // CHECK-NOT: subview
  // CHECK-NOT: expand_shape
  // CHECK-NOT: collapse_shape

  %base = amdgpu.make_dma_base %mem[%global_i, %global_j], %lds[%lds_i, %lds_j]
    : memref<64x128xf16, #gpu_global_addrspace>, memref<64x64xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_base<f16>
  func.return %base : !amdgpu.tdm_base<f16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_make_gather_dma_base_nop
// CHECK-SAME: %[[MEM:.*]]: memref<64x128xf16, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xf16, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index
func.func @test_make_gather_dma_base_nop(%mem: memref<64x128xf16, #gpu_global_addrspace>, %lds: memref<64x64xf16, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index) -> !amdgpu.tdm_gather_base<f16, i16> {
  // CHECK: %[[BASE:.*]] = amdgpu.make_gather_dma_base %[[MEM]][%[[GLOBAL_I]], %[[GLOBAL_J]]], %[[LDS]][%[[LDS_I]], %[[LDS_J]]]
  // CHECK-SAME: memref<64x128xf16, #gpu.address_space<global>>, memref<64x64xf16, #gpu.address_space<workgroup>> -> !amdgpu.tdm_gather_base<f16, i16>
  // CHECK: return %[[BASE]] : !amdgpu.tdm_gather_base<f16, i16>
  // CHECK-NOT: subview
  // CHECK-NOT: expand_shape
  // CHECK-NOT: collapse_shape

  %base = amdgpu.make_gather_dma_base %mem[%global_i, %global_j], %lds[%lds_i, %lds_j]
    : memref<64x128xf16, #gpu_global_addrspace>, memref<64x64xf16, #gpu_lds_addrspace> -> !amdgpu.tdm_gather_base<f16, i16>
  func.return %base : !amdgpu.tdm_gather_base<f16, i16>
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: #[[GLOBAL_ASYNC_MAP:.*]] = affine_map<()[s0] -> (s0 + 32)>
// CHECK: #[[GLOBAL_ASYNC_MAP1:.*]] = affine_map<()[s0] -> (s0 + 64)>

// CHECK: func @test_global_load_async_to_lds_both_fold_masked
// CHECK-SAME: %[[SRC:.*]]: memref<64x128xf32, #gpu.address_space<global>>, %[[LDS:.*]]: memref<4096xf32, #gpu.address_space<workgroup>>, %[[GLOBAL_I:.*]]: index, %[[GLOBAL_J:.*]]: index, %[[LDS_I:.*]]: index, %[[LDS_J:.*]]: index, %[[MASK:.*]]: i1
func.func @test_global_load_async_to_lds_both_fold_masked(%src: memref<64x128xf32, #gpu_global_addrspace>, %lds: memref<4096xf32, #gpu_lds_addrspace>, %global_i: index, %global_j: index, %lds_i: index, %lds_j: index, %mask: i1) {
  // CHECK: %[[GI:.*]] = affine.apply #[[GLOBAL_ASYNC_MAP]]()[%[[GLOBAL_I]]]
  // CHECK: %[[GJ:.*]] = affine.apply #[[GLOBAL_ASYNC_MAP1]]()[%[[GLOBAL_J]]]
  // CHECK: %[[LDS_IDX:.*]] = affine.linearize_index [%[[LDS_I]], %[[LDS_J]]] by (64, 64) : index
  // CHECK: amdgpu.global_load_async_to_lds %[[SRC]][%[[GI]], %[[GJ]]], %[[LDS]][%[[LDS_IDX]]], %[[MASK]]
  // CHECK-SAME: vector<4xf32>, memref<64x128xf32, #gpu.address_space<global>>, memref<4096xf32, #gpu.address_space<workgroup>>

  %subview = memref.subview %src[32, 64][32, 64][1, 1]
    : memref<64x128xf32, #gpu_global_addrspace>
    to memref<32x64xf32, strided<[128, 1], offset: 4160>, #gpu_global_addrspace>
  %expand_lds = memref.expand_shape %lds [[0, 1]] output_shape [64, 64]
    : memref<4096xf32, #gpu_lds_addrspace>
    into memref<64x64xf32, #gpu_lds_addrspace>
  amdgpu.global_load_async_to_lds %subview[%global_i, %global_j], %expand_lds[%lds_i, %lds_j], %mask
    : vector<4xf32>, memref<32x64xf32, strided<[128, 1], offset: 4160>, #gpu_global_addrspace>,
      memref<64x64xf32, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_global_load_async_to_lds_no_mask_dst_collapse
// CHECK-SAME: %[[SRC:.*]]: memref<8192xi32, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xi32, #gpu.address_space<workgroup>>, %[[SRC_IDX:.*]]: index, %[[DST_IDX:.*]]: index
func.func @test_global_load_async_to_lds_no_mask_dst_collapse(%src: memref<8192xi32, #gpu_global_addrspace>, %lds: memref<64x64xi32, #gpu_lds_addrspace>, %src_idx: index, %dst_idx: index) {
  // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[DST_IDX]] into (64, 64) : index, index
  // CHECK: amdgpu.global_load_async_to_lds %[[SRC]][%[[SRC_IDX]]], %[[LDS]][%[[INDICES]]#0, %[[INDICES]]#1] :
  // CHECK-SAME: i32, memref<8192xi32, #gpu.address_space<global>>, memref<64x64xi32, #gpu.address_space<workgroup>>

  %collapse_lds = memref.collapse_shape %lds [[0, 1]]
    : memref<64x64xi32, #gpu_lds_addrspace>
    into memref<4096xi32, #gpu_lds_addrspace>
  amdgpu.global_load_async_to_lds %src[%src_idx], %collapse_lds[%dst_idx]
    : i32, memref<8192xi32, #gpu_global_addrspace>,
      memref<4096xi32, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_lds_addrspace = #gpu.address_space<workgroup>
#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_global_load_async_to_lds_masked_dst_collapse
// CHECK-SAME: %[[SRC:.*]]: memref<8192xi32, #gpu.address_space<global>>, %[[LDS:.*]]: memref<64x64xi32, #gpu.address_space<workgroup>>, %[[SRC_IDX:.*]]: index, %[[DST_IDX:.*]]: index, %[[MASK:.*]]: i1
func.func @test_global_load_async_to_lds_masked_dst_collapse(%src: memref<8192xi32, #gpu_global_addrspace>, %lds: memref<64x64xi32, #gpu_lds_addrspace>, %src_idx: index, %dst_idx: index, %mask: i1) {
  // CHECK: %[[INDICES:.*]]:2 = affine.delinearize_index %[[DST_IDX]] into (64) : index, index
  // CHECK: amdgpu.global_load_async_to_lds %[[SRC]][%[[SRC_IDX]]], %[[LDS]][%[[INDICES]]#0, %[[INDICES]]#1], %[[MASK]] :
  // CHECK-SAME: i32, memref<8192xi32, #gpu.address_space<global>>, memref<64x64xi32, #gpu.address_space<workgroup>>

  %collapse_lds = memref.collapse_shape %lds [[0, 1]]
    : memref<64x64xi32, #gpu_lds_addrspace>
    into memref<4096xi32, #gpu_lds_addrspace>
  amdgpu.global_load_async_to_lds %src[%src_idx], %collapse_lds[%dst_idx], %mask
    : i32, memref<8192xi32, #gpu_global_addrspace>,
      memref<4096xi32, #gpu_lds_addrspace>
  func.return
}

// -----

#gpu_global_addrspace = #gpu.address_space<global>

// CHECK: func @test_global_transpose_load_expand_shape
// CHECK-SAME: %[[SRC:.*]]: memref<8192xf16, #gpu.address_space<global>>, %[[I:.*]]: index, %[[J:.*]]: index
func.func @test_global_transpose_load_expand_shape(%src: memref<8192xf16, #gpu_global_addrspace>, %i: index, %j: index) -> vector<8xf16> {
  // CHECK: %[[IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (64, 128) : index
  // CHECK: amdgpu.global_transpose_load %[[SRC]][%[[IDX]]]
  // CHECK-SAME: memref<8192xf16, #gpu.address_space<global>> -> vector<8xf16>

  %expand_src = memref.expand_shape %src [[0, 1]] output_shape [64, 128]
    : memref<8192xf16, #gpu_global_addrspace>
    into memref<64x128xf16, #gpu_global_addrspace>
  %result = amdgpu.global_transpose_load %expand_src[%i, %j]
    : memref<64x128xf16, #gpu_global_addrspace> -> vector<8xf16>
  return %result : vector<8xf16>
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_ds_barrier_ops_expand_shape
// CHECK-SAME: %[[BARRIER:.*]]: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, %[[I:.*]]: index, %[[J:.*]]: index, %[[PARTICIPANTS:.*]]: i32, %[[COUNT:.*]]: i64
func.func @test_ds_barrier_ops_expand_shape(%barrier: memref<16x!amdgpu.ds_barrier_state, #gpu_wg>, %i: index, %j: index, %participants: i32, %count: i64) -> !amdgpu.ds_barrier_state {
  // CHECK: %[[INIT_IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (4, 4) : index
  // CHECK: amdgpu.ds_barrier_init %[[BARRIER]][%[[INIT_IDX]]], %[[PARTICIPANTS]]
  // CHECK-SAME: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, i32
  // CHECK: %[[POLL_IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (4, 4) : index
  // CHECK: %[[STATE:.*]] = amdgpu.ds_barrier_poll_state %[[BARRIER]][%[[POLL_IDX]]]
  // CHECK-SAME: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>> -> !amdgpu.ds_barrier_state
  // CHECK: %[[ASYNC_IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (4, 4) : index
  // CHECK: amdgpu.ds_async_barrier_arrive %[[BARRIER]][%[[ASYNC_IDX]]]
  // CHECK-SAME: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>
  // CHECK: %[[ARRIVE_IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (4, 4) : index
  // CHECK: %[[OLD:.*]] = amdgpu.ds_barrier_arrive %[[BARRIER]][%[[ARRIVE_IDX]]], %[[COUNT]]
  // CHECK-SAME: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, i64 -> !amdgpu.ds_barrier_state

  %expand_barrier = memref.expand_shape %barrier [[0, 1]] output_shape [4, 4]
    : memref<16x!amdgpu.ds_barrier_state, #gpu_wg>
    into memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>
  amdgpu.ds_barrier_init %expand_barrier[%i, %j], %participants
    : memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>, i32
  %state = amdgpu.ds_barrier_poll_state %expand_barrier[%i, %j]
    : memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg> -> !amdgpu.ds_barrier_state
  amdgpu.ds_async_barrier_arrive %expand_barrier[%i, %j]
    : memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>
  %old_state = amdgpu.ds_barrier_arrive %expand_barrier[%i, %j], %count
    : memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>, i64 -> !amdgpu.ds_barrier_state
  return %old_state : !amdgpu.ds_barrier_state
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_make_dma_descriptor_atomic_barrier_expand_shape
// CHECK-SAME: %[[BASE:.*]]: !amdgpu.tdm_base<i32>, %[[BARRIER:.*]]: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, %[[I:.*]]: index, %[[J:.*]]: index
func.func @test_make_dma_descriptor_atomic_barrier_expand_shape(%base: !amdgpu.tdm_base<i32>, %barrier: memref<16x!amdgpu.ds_barrier_state, #gpu_wg>, %i: index, %j: index) -> !amdgpu.tdm_descriptor {
  // CHECK: %[[BARRIER_IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (4, 4) : index
  // CHECK: %[[DESC:.*]] = amdgpu.make_dma_descriptor %[[BASE]]
  // CHECK-SAME: globalSize [64, 64]
  // CHECK-SAME: globalStride [64, 1]
  // CHECK-SAME: sharedSize [64, 64]
  // CHECK-SAME: atomicBarrier(%[[BARRIER]][%[[BARRIER_IDX]]] : memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>)

  %expand_barrier = memref.expand_shape %barrier [[0, 1]] output_shape [4, 4]
    : memref<16x!amdgpu.ds_barrier_state, #gpu_wg>
    into memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>
  %desc = amdgpu.make_dma_descriptor %base
    globalSize [64, 64]
    globalStride [64, 1]
    sharedSize [64, 64]
    atomicBarrier(%expand_barrier[%i, %j] : memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>)
    : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor
  return %desc : !amdgpu.tdm_descriptor
}

// -----

#gpu_wg = #gpu.address_space<workgroup>

// CHECK: func @test_make_gather_dma_descriptor_atomic_barrier_expand_shape
// CHECK-SAME: %[[BASE:.*]]: !amdgpu.tdm_gather_base<i32, i32>, %[[INDICES:.*]]: vector<8xi32>, %[[BARRIER:.*]]: memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, %[[I:.*]]: index, %[[J:.*]]: index
func.func @test_make_gather_dma_descriptor_atomic_barrier_expand_shape(%base: !amdgpu.tdm_gather_base<i32, i32>, %indices: vector<8xi32>, %barrier: memref<16x!amdgpu.ds_barrier_state, #gpu_wg>, %i: index, %j: index) -> !amdgpu.tdm_descriptor {
  // CHECK: %[[BARRIER_IDX:.*]] = affine.linearize_index disjoint [%[[I]], %[[J]]] by (4, 4) : index
  // CHECK: %[[DESC:.*]] = amdgpu.make_gather_dma_descriptor %[[BASE]][%[[INDICES]]]
  // CHECK-SAME: globalSize [64, 64]
  // CHECK-SAME: globalStride [64, 1]
  // CHECK-SAME: sharedSize [64, 64]
  // CHECK-SAME: atomicBarrier(%[[BARRIER]][%[[BARRIER_IDX]]] : memref<16x!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>)

  %expand_barrier = memref.expand_shape %barrier [[0, 1]] output_shape [4, 4]
    : memref<16x!amdgpu.ds_barrier_state, #gpu_wg>
    into memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>
  %desc = amdgpu.make_gather_dma_descriptor %base[%indices]
    globalSize [64, 64]
    globalStride [64, 1]
    sharedSize [64, 64]
    atomicBarrier(%expand_barrier[%i, %j] : memref<4x4x!amdgpu.ds_barrier_state, #gpu_wg>)
    : !amdgpu.tdm_gather_base<i32, i32>, vector<8xi32> -> !amdgpu.tdm_descriptor
  return %desc : !amdgpu.tdm_descriptor
}

// -----

// CHECK: func @test_make_dma_descriptor_no_atomic_barrier
// CHECK-SAME: %[[BASE:.*]]: !amdgpu.tdm_base<i32>
func.func @test_make_dma_descriptor_no_atomic_barrier(%base: !amdgpu.tdm_base<i32>) -> !amdgpu.tdm_descriptor {
  // CHECK: %[[DESC:.*]] = amdgpu.make_dma_descriptor %[[BASE]]
  // CHECK-SAME: globalSize [64, 64]
  // CHECK-SAME: globalStride [64, 1]
  // CHECK-SAME: sharedSize [64, 64]
  // CHECK-NOT: atomicBarrier

  %desc = amdgpu.make_dma_descriptor %base
    globalSize [64, 64]
    globalStride [64, 1]
    sharedSize [64, 64]
    : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor
  return %desc : !amdgpu.tdm_descriptor
}