184 lines · plain
1// RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx942 | FileCheck %s2// RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx950 | FileCheck %s3 4#gpu_global_addrspace = 15#gpu_lds_addrspace = 36#amdgpu_fat_buffer_addrspace = 77 8// CHECK-LABEL: func @global_load_to_rocdl_f329// CHECK-SAME: (%[[ARG0:.*]]: memref<128x72xf32, 1>)10func.func @global_load_to_rocdl_f32(%global : memref<128x72xf32, #gpu_global_addrspace>) {11 %c0 = arith.constant 0 : index12 %c12 = arith.constant 12 : index13 %c32 = arith.constant 32 : index14 %alloc = memref.alloc() : memref<64x64xf32, #gpu_lds_addrspace>15 // CHECK: %[[GLOBAL_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]16 17 // CHECK: %[[C0:.*]] = arith.constant 0 : index18 // CHECK: %[[IC0:.*]] = builtin.unrealized_conversion_cast %c0 : index to i6419 // CHECK: %[[C12:.*]] = arith.constant 12 : index20 // CHECK: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]21 // CHECK: %[[C32:.*]] = arith.constant 32 : index22 // CHECK: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]23 24 // CHECK: %[[ALLOC:.*]] = memref.alloc()25 // CHECK: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast26 // CHECK: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[GLOBAL_DESC]][1]27 28 // CHECK: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i6429 // CHECK: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i6430 // CHECK: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i6431 32 // CHECK: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]33 // CHECK: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]34 35 // CHECK: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i6436 // CHECK: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C64]] : i6437 // CHECK: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i6438 39 // CHECK: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]40 // CHECK: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 441 amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]42 : f32, memref<128x72xf32, #gpu_global_addrspace>, memref<64x64xf32, #gpu_lds_addrspace>43 func.return44}45 46// CHECK-LABEL: func @global_load_to_rocdl_i847// CHECK-SAME: (%[[ARG0:.*]]: memref<128x72xi8, 1>)48func.func @global_load_to_rocdl_i8(%global : memref<128x72xi8, #gpu_global_addrspace>) {49 // CHECK: %[[GLOBAL_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]50 51 // CHECK: %[[C0:.*]] = arith.constant 0 : index52 // CHECK: %[[IC0:.*]] = builtin.unrealized_conversion_cast %c0 : index to i6453 // CHECK: %[[C12:.*]] = arith.constant 12 : index54 // CHECK: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]55 // CHECK: %[[C32:.*]] = arith.constant 32 : index56 // CHECK: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]57 58 // CHECK: %[[ALLOC:.*]] = memref.alloc()59 // CHECK: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast %[[ALLOC]]60 // CHECK: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[GLOBAL_DESC]][1]61 62 // CHECK: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i6463 // CHECK: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i6464 // CHECK: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i6465 66 // CHECK: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]67 // CHECK: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]68 69 // CHECK: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i6470 // CHECK: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C64]] : i6471 // CHECK: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i6472 73 // CHECK: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]74 // CHECK: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 175 %c0 = arith.constant 0 : index76 %c12 = arith.constant 12 : index77 %c32 = arith.constant 32 : index78 %alloc = memref.alloc() : memref<64x64xi8, #gpu_lds_addrspace>79 amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]80 : i8, memref<128x72xi8, #gpu_global_addrspace>, memref<64x64xi8, #gpu_lds_addrspace>81 func.return82}83 84// CHECK-LABEL: func @global_load_to_rocdl_vec85// CHECK-SAME: (%[[ARG0:.*]]: memref<128x72xi16, 1>)86func.func @global_load_to_rocdl_vec(%global : memref<128x72xi16, #gpu_global_addrspace>) {87 // CHECK: %[[GLOBAL_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]88 89 // CHECK: %[[C0:.*]] = arith.constant 0 : index90 // CHECK: %[[IC0:.*]] = builtin.unrealized_conversion_cast %c0 : index to i6491 // CHECK: %[[C12:.*]] = arith.constant 12 : index92 // CHECK: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]93 // CHECK: %[[C32:.*]] = arith.constant 32 : index94 // CHECK: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]95 96 // CHECK: %[[ALLOC:.*]] = memref.alloc()97 // CHECK: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast %[[ALLOC]]98 // CHECK: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[GLOBAL_DESC]][1]99 100 // CHECK: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i64101 // CHECK: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i64102 // CHECK: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i64103 104 // CHECK: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]105 // CHECK: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]106 107 // CHECK: %[[C128:.*]] = llvm.mlir.constant(128 : index) : i64108 // CHECK: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C128]] : i64109 // CHECK: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i64110 111 // CHECK: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]112 // CHECK: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 4113 %c0 = arith.constant 0 : index114 %c12 = arith.constant 12 : index115 %c32 = arith.constant 32 : index116 %alloc = memref.alloc() : memref<64x128xi16, #gpu_lds_addrspace>117 amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]118 : vector<2 x i16>, memref<128x72xi16, #gpu_global_addrspace>, memref<64x128xi16, #gpu_lds_addrspace>119 func.return120}121 122// CHECK-LABEL: func @global_load_to_rocdl_dynamic_indices123// CHECK-SAME: (%[[ARG0:.*]]: memref<512xi32, 1>, %[[SRC_IDX:.*]]: index, %[[DST_IDX:.*]]: index)124func.func @global_load_to_rocdl_dynamic_indices(%global : memref<512xi32, #gpu_global_addrspace>, %src_idx : index, %dst_idx : index) {125 // CHECK: %[[DSTIDX_CAST:.*]] = builtin.unrealized_conversion_cast %[[DST_IDX]]126 // CHECK: %[[SRCIDX_CAST:.*]] = builtin.unrealized_conversion_cast %[[SRC_IDX]]127 // CHECK: %[[GLOBAL_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]128 // CHECK: %[[ALLOC:.*]] = memref.alloc()129 // CHECK: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast %[[ALLOC]]130 // CHECK: %[[C0:.*]] = arith.constant 0 : index131 // CHECK: %[[C0_I64:.*]] = builtin.unrealized_conversion_cast %[[C0]] : index to i64132 // CHECK: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[GLOBAL_DESC]][1]133 // CHECK: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRCIDX_CAST]]]134 // CHECK: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]135 // CHECK: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i64136 // CHECK: %[[DSTIDX:.*]] = llvm.mul %[[DSTIDX_CAST]], %[[C64]] : i64137 // CHECK: %[[DSTIDX1:.*]] = llvm.add %[[DSTIDX]], %[[C0_I64]] : i64138 // CHECK: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DSTIDX1]]]139 // CHECK: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 4140 %alloc = memref.alloc() : memref<4x64xi32, #gpu_lds_addrspace>141 %c0 = arith.constant 0 : index142 amdgpu.gather_to_lds %global[%src_idx], %alloc[%dst_idx, %c0]143 : i32, memref<512xi32, #gpu_global_addrspace>, memref<4x64xi32, #gpu_lds_addrspace>144 func.return145}146 147// CHECK-LABEL: func @fat_buffer_load_to_rocdl_f32148// CHECK-SAME: (%[[ARG0:.*]]: memref<128x72xf32, 7>)149func.func @fat_buffer_load_to_rocdl_f32(%global : memref<128x72xf32, #amdgpu_fat_buffer_addrspace>) {150 %c0 = arith.constant 0 : index151 %c12 = arith.constant 12 : index152 %c32 = arith.constant 32 : index153 %alloc = memref.alloc() : memref<64x64xf32, #gpu_lds_addrspace>154 // CHECK: %[[BUFFER_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]155 156 // CHECK: %[[C0:.*]] = arith.constant 0 : index157 // CHECK: %[[IC0:.*]] = builtin.unrealized_conversion_cast %[[C0]] : index to i64158 // CHECK: %[[C12:.*]] = arith.constant 12 : index159 // CHECK: %[[IC12:.*]] = builtin.unrealized_conversion_cast %[[C12]]160 // CHECK: %[[C32:.*]] = arith.constant 32 : index161 // CHECK: %[[IC32:.*]] = builtin.unrealized_conversion_cast %[[C32]]162 163 // CHECK: %[[ALLOC:.*]] = memref.alloc()164 // CHECK: %[[LDS_DESC:.*]] = builtin.unrealized_conversion_cast165 // CHECK: %[[GLOBAL_BASE:.*]] = llvm.extractvalue %[[BUFFER_DESC]][1]166 167 // CHECK: %[[C72:.*]] = llvm.mlir.constant(72 : index) : i64168 // CHECK: %[[MUL:.*]] = llvm.mul %[[IC12]], %[[C72]] : i64169 // CHECK: %[[SRC_OFFSET:.*]] = llvm.add %[[MUL]], %[[IC0]] : i64170 171 // CHECK: %[[GLOBAL_PTR:.*]] = llvm.getelementptr %[[GLOBAL_BASE]][%[[SRC_OFFSET]]]172 // CHECK: %[[LDS_BASE:.*]] = llvm.extractvalue %[[LDS_DESC]][1]173 174 // CHECK: %[[C64:.*]] = llvm.mlir.constant(64 : index) : i64175 // CHECK: %[[MUL_2:.*]] = llvm.mul %[[IC32]], %[[C64]] : i64176 // CHECK: %[[DST_OFFSET:.*]] = llvm.add %[[MUL_2]], %[[IC0]] : i64177 178 // CHECK: %[[LDS_PTR:.*]] = llvm.getelementptr %[[LDS_BASE]][%[[DST_OFFSET]]]179 // CHECK: rocdl.load.to.lds %[[GLOBAL_PTR]], %[[LDS_PTR]], 4180 amdgpu.gather_to_lds %global[%c12, %c0], %alloc[%c32, %c0]181 : f32, memref<128x72xf32, #amdgpu_fat_buffer_addrspace>, memref<64x64xf32, #gpu_lds_addrspace>182 func.return183}184