164 lines · plain
1// RUN: mlir-opt --amdgpu-fold-memrefs-ops --split-input-file %s | FileCheck %s2 3#gpu_lds_addrspace = 34 5// CHECK: func @test_subview_folding6// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index7func.func @test_subview_folding(%offset_i: index, %offset_j: index) {8 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, 3>9 // CHECK: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>10 // CHECK: %[[C0:.*]] = arith.constant 0 : index11 // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[ARG0]], %[[ARG1]]], %[[LOCAL]][%[[C0]], %[[C0]]]12 // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, 3>13 14 %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>15 %mem = memref.alloc() : memref<64x128xf16>16 %subview = memref.subview %mem[0, 0][32, 64][1, 1] : memref<64x128xf16> to memref<32x64xf16, strided<[128, 1]>>17 %c0 = arith.constant 0 : index18 amdgpu.gather_to_lds %subview[%offset_i, %offset_j], %alloc[%c0, %c0]19 : vector<8xf16>, memref<32x64xf16, strided<[128, 1]>>, memref<64x64xf16, #gpu_lds_addrspace>20 func.return21}22 23// -----24 25#gpu_lds_addrspace = 326 27// CHECK: #[[MAP:.*]] = affine_map<()[s0] -> (s0 + 32)>28// CHECK: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 64)>29 30// CHECK: func @subview_folding_offset31// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index32func.func @subview_folding_offset(%offset_i: index, %offset_j: index) {33 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, 3>34 // CHECK: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>35 // CHECK: %[[C0:.*]] = arith.constant 0 : index36 // CHECK: %[[IDX0:.*]] = affine.apply #[[MAP]]()[%[[ARG0]]]37 // CHECK: %[[IDX1:.*]] = affine.apply #[[MAP1]]()[%[[ARG1]]]38 // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[IDX0]], %[[IDX1]]], %[[LOCAL]][%[[C0]], %[[C0]]]39 // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, 3>40 41 %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>42 %mem = memref.alloc() : memref<64x128xf16>43 %subview = memref.subview %mem[32, 64][32, 64][1, 1] : memref<64x128xf16> to memref<32x64xf16, strided<[128, 1], offset: 4160>>44 %c0 = arith.constant 0 : index45 amdgpu.gather_to_lds %subview[%offset_i, %offset_j], %alloc[%c0, %c0]46 : vector<8xf16>, memref<32x64xf16, strided<[128, 1], offset: 4160>>, memref<64x64xf16, #gpu_lds_addrspace>47 func.return48}49 50// -----51 52#gpu_lds_addrspace = 353 54// CHECK: func @test_expand_shape55// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index56func.func @test_expand_shape(%offset_i: index, %offset_j: index) {57 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, 3>58 // CHECK: %[[MEM:.*]] = memref.alloc() : memref<8192xf16>59 // CHECK: %[[C0:.*]] = arith.constant 0 : index60 // CHECK: %[[IDXM:.*]] = affine.linearize_index [%[[ARG0]], %[[ARG1]]] by (64, 128) : index61 // CHECK: %[[IDXL:.*]] = affine.linearize_index [%[[C0]], %[[C0]]] by (64, 64) : index62 // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[IDXM]]], %[[LOCAL]][%[[IDXL]]]63 // CHECK-SAME: vector<8xf16>, memref<8192xf16>, memref<4096xf16, 3>64 65 %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>66 %mem = memref.alloc() : memref<8192xf16>67 %expand_mem = memref.expand_shape %mem [[0, 1]] output_shape [64, 128] : memref<8192xf16> into memref<64x128xf16>68 %expand_alloc = memref.expand_shape %alloc [[0, 1]] output_shape [64, 64] : memref<4096xf16, #gpu_lds_addrspace> into memref<64x64xf16, #gpu_lds_addrspace>69 %c0 = arith.constant 0 : index70 amdgpu.gather_to_lds %expand_mem[%offset_i, %offset_j], %expand_alloc[%c0, %c0]71 : vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, #gpu_lds_addrspace>72 func.return73}74 75// -----76 77#gpu_lds_addrspace = 378 79// CHECK: func @test_collapse_shape80// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index81func.func @test_collapse_shape(%offset_i: index, %offset_j: index) {82 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<64x64xf16, 3>83 // CHECK: %[[MEM:.*]] = memref.alloc() : memref<64x128xf16>84 // CHECK: %[[C0:.*]] = arith.constant 0 : index85 // CHECK: %[[INDICES_MEM:.*]]:2 = affine.delinearize_index %[[ARG0]] into (64, 128) : index, index86 // CHECK: %[[INDICES_LDS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (64, 64) : index, index87 // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[INDICES_MEM]]#0, %[[INDICES_MEM]]#1], %[[LOCAL]][%[[INDICES_LDS]]#0, %[[INDICES_LDS]]#1]88 // CHECK-SAME: vector<8xf16>, memref<64x128xf16>, memref<64x64xf16, 3>89 90 %alloc = memref.alloc() : memref<64x64xf16, #gpu_lds_addrspace>91 %collapse_alloc = memref.collapse_shape %alloc [[0, 1]] : memref<64x64xf16, #gpu_lds_addrspace> into memref<4096xf16, #gpu_lds_addrspace>92 %mem = memref.alloc() : memref<64x128xf16>93 %collapse_mem = memref.collapse_shape %mem [[0, 1]] : memref<64x128xf16> into memref<8192xf16>94 %c0 = arith.constant 0 : index95 amdgpu.gather_to_lds %collapse_mem[%offset_i], %collapse_alloc[%offset_j]96 : vector<8xf16>, memref<8192xf16>, memref<4096xf16, #gpu_lds_addrspace>97 func.return98}99 100 101// -----102 103#gpu_lds_addrspace = 3104 105 106// CHECK: func @test_expand_shape_src_raw_buffer107// CHECK-SAME: %[[ARG0:.*]]: memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index108func.func @test_expand_shape_src_raw_buffer(%mem : memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %offset_i: index, %offset_j: index) {109 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, 3>110 // CHECK: %[[C0:.*]] = arith.constant 0 : index111 // CHECK: %[[IDXM:.*]] = affine.linearize_index [%[[ARG1]], %[[ARG2]]] by (64, 128) : index112 // CHECK: amdgpu.gather_to_lds %[[ARG0]][%[[IDXM]]], %[[LOCAL]][%[[C0]]]113 // CHECK-SAME: vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, 3>114 115 %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>116 %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>>117 118 %c0 = arith.constant 0 : index119 amdgpu.gather_to_lds %expand_mem[%offset_i, %offset_j], %alloc[%c0]120 : vector<8xf16>, memref<64x128xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu_lds_addrspace>121 func.return122}123 124// -----125 126#gpu_lds_addrspace = 3127 128// CHECK: func @test_expand_shape_dst_only129// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index130func.func @test_expand_shape_dst_only(%offset_i: index, %offset_j: index) {131 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, 3>132 // CHECK: %[[MEM:.*]] = memref.alloc() : memref<8192xf16>133 // CHECK: %[[C0:.*]] = arith.constant 0 : index134 // CHECK: %[[IDX_LDS:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (64, 64) : index135 // CHECK: amdgpu.gather_to_lds %[[MEM]][%[[ARG0]]], %[[LOCAL]][%[[IDX_LDS]]]136 // CHECK-SAME: vector<8xf16>, memref<8192xf16>, memref<4096xf16, 3>137 138 %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>139 %mem = memref.alloc() : memref<8192xf16>140 %expand_alloc = memref.expand_shape %alloc [[0, 1]] output_shape [64, 64] : memref<4096xf16, #gpu_lds_addrspace> into memref<64x64xf16, #gpu_lds_addrspace>141 142 %c0 = arith.constant 0 : index143 amdgpu.gather_to_lds %mem[%offset_i], %expand_alloc[%offset_j, %c0]144 : vector<8xf16>, memref<8192xf16>, memref<64x64xf16, #gpu_lds_addrspace>145 func.return146}147 148// -----149 150#gpu_lds_addrspace = 3151 152// CHECK: func @test_nop153// CHECK-SAME: %[[ARG0:.*]]: memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index154func.func @test_nop(%mem : memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, %offset_i: index, %offset_j: index) {155 // CHECK: %[[LOCAL:.*]] = memref.alloc() : memref<4096xf16, 3>156 // CHECK: amdgpu.gather_to_lds %[[ARG0]][%[[ARG1]]], %[[LOCAL]][%[[ARG2]]]157 // CHECK-SAME: vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, 3>158 159 %alloc = memref.alloc() : memref<4096xf16, #gpu_lds_addrspace>160 amdgpu.gather_to_lds %mem[%offset_i], %alloc[%offset_j]161 : vector<8xf16>, memref<8192xf16, #amdgpu.address_space<fat_raw_buffer>>, memref<4096xf16, #gpu_lds_addrspace>162 func.return163}164