170 lines · plain
1// RUN: mlir-opt %s --amdgpu-maskedload-to-load --split-input-file | FileCheck %s2 3// CHECK-LABEL: func @transfer_to_maskedload_fatrawbuffer(4// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>5// CHECK-SAME: %[[ARG1:.*]]: index6// CHECK-SAME: %[[ARG2:.*]]: vector<4xi1>7// CHECK-SAME: %[[ARG3:.*]]: vector<4xf32>8func.func @transfer_to_maskedload_fatrawbuffer(%mem : memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>, %idx : index, %mask : vector<4xi1>, %passthru : vector<4xf32>) -> vector<4xf32> {9 %res = vector.maskedload %mem[%idx, %idx], %mask, %passthru : memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>, vector<4xi1>, vector<4xf32> into vector<4xf32>10 return %res : vector<4xf32>11}12 13// CHECK: %[[IF:.*]] = scf.if14// CHECK: vector.maskedload %[[ARG0]][%[[ARG1]], %[[ARG1]]]15 16// CHECK: } else {17// CHECK: %[[LOAD:.*]] = vector.load %arg0[%arg1, %arg1]18// CHECK: %[[SELECT:.*]] = arith.select %[[ARG2]], %[[LOAD]]19 20// CHECK: return %[[IF]] : vector<4xf32>21 22// -----23 24// CHECK: #map = affine_map<()[s0, s1] -> (s0 * 8 + s1)>25// CHECK-LABEL: func @transfer_to_maskedload_fatrawbuffer_f16(26// CHECK-SAME: %[[ARG0:.+]]: memref<8x8xf16, #amdgpu.address_space<fat_raw_buffer>>,27// CHECK-SAME: %[[ARG1:.+]]: index, %[[ARG2:.+]]: index,28// CHECK-SAME: %[[ARG3:.+]]: vector<4xi1>29// CHECK-SAME: %[[ARG4:.+]]: vector<4xf16>30func.func @transfer_to_maskedload_fatrawbuffer_f16(%mem : memref<8x8xf16, #amdgpu.address_space<fat_raw_buffer>>, %idx0 : index, %idx1 : index, %mask : vector<4xi1>, %passthru : vector<4xf16>) -> vector<4xf16> {31 %res = vector.maskedload %mem[%idx0, %idx1], %mask, %passthru : memref<8x8xf16, #amdgpu.address_space<fat_raw_buffer>>, vector<4xi1>, vector<4xf16> into vector<4xf16>32 return %res : vector<4xf16>33}34// CHECK-DAG: %[[C0:.*]] = arith.constant 035// CHECK-DAG: %[[SIZE:.*]] = arith.constant 6436// CHECK-DAG: %[[BYTES:.*]] = arith.constant 237// CHECK-DAG: %[[C4:.*]] = arith.constant 438 39// CHECK: %[[LINEAR:.*]] = affine.apply #map()[%[[ARG1]], %[[ARG2]]]40// CHECK: %[[DELTA:.*]] = arith.subi %[[SIZE]], %[[LINEAR]]41// CHECK: %[[COND1:.*]] = arith.cmpi ult, %[[DELTA]], %[[C4]]42 43// CHECK: %[[REM:.*]] = arith.remui %[[DELTA]], %[[BYTES]]44// CHECK: %[[COND2:.*]] = arith.cmpi ne, %[[REM]], %[[C0]]45 46// CHECK: %[[COND:.*]] = arith.andi %[[COND1]], %[[COND2]]47// CHECK: %[[IF:.*]] = scf.if %[[COND]] -> (vector<4xf16>) {48// CHECK: vector.maskedload %[[ARG0]][%[[ARG1]], %[[ARG2]]]49// CHECK: } else {50// CHECK: %[[LOAD:.*]] = vector.load %[[ARG0]][%[[ARG1]], %[[ARG2]]]51// CHECK: return %[[IF]] : vector<4xf16>52 53// -----54 55// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 * s1 + s2)>56// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (s0 * s1, s2)>57// CHECK: func @transfer_to_maskedload_fatrawbuffer_dynamic_i8(58// CHECK-SAME: %[[ARG0:.*]]: memref<?x?xi8, #amdgpu.address_space<fat_raw_buffer>>59// CHECK-SAME: %[[ARG1:.*]]: index, %[[ARG2:.*]]: index60// CHECK-SAME: %[[ARG3:.*]]: vector<4xi1>61// CHECK-SAME: %[[ARG4:.*]]: vector<4xi8>62func.func @transfer_to_maskedload_fatrawbuffer_dynamic_i8(%mem : memref<?x?xi8, #amdgpu.address_space<fat_raw_buffer>>, %idx0 : index, %idx1 : index, %mask : vector<4xi1>, %passthru : vector<4xi8>) -> vector<4xi8> {63 %res = vector.maskedload %mem[%idx0, %idx1], %mask, %passthru : memref<?x?xi8, #amdgpu.address_space<fat_raw_buffer>>, vector<4xi1>, vector<4xi8> into vector<4xi8>64 return %res : vector<4xi8>65}66// CHECK: %[[C0:.*]] = arith.constant 0 : index67// CHECK: %[[C4:.*]] = arith.constant 4 : index68// CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG0]]69// CHECK-DAG: %[[SIZE:.*]] = affine.max #[[MAP1]]()[%[[STRIDES]]#0, %[[SIZES]]#0, %[[SIZES]]#1]70// CHECK-DAG: %[[LINEAR:.*]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[STRIDES]]#0, %[[ARG2]]]71// CHECK: %[[IF:.*]] = scf.if72// CHECK: return73 74// -----75 76// CHECK-LABEL: func @transfer_to_maskedload_regular(77// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32>78// CHECK-SAME: %[[ARG1:.*]]: index79// CHECK-SAME: %[[ARG2:.*]]: vector<4xi1>80// CHECK-SAME: %[[ARG3:.*]]: vector<4xf32>81func.func @transfer_to_maskedload_regular(%mem : memref<8x8xf32>, %idx : index, %mask : vector<4xi1>, %passthru : vector<4xf32>) -> vector<4xf32> {82 %res = vector.maskedload %mem[%idx, %idx], %mask, %passthru : memref<8x8xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>83 return %res : vector<4xf32>84}85// CHECK: %[[RES:.*]] = vector.maskedload %[[ARG0]][%[[ARG1]], %[[ARG1]]], %[[ARG2]], %[[ARG3]]86// CHECK: return %[[RES]] : vector<4xf32>87 88// -----89 90// CHECK-LABEL: func @transfer_to_maskedload_addrspace(91// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32, #gpu.address_space<workgroup>>92// CHECK-SAME: %[[ARG1:.*]]: index93// CHECK-SAME: %[[ARG2:.*]]: vector<4xi1>94// CHECK-SAME: %[[ARG3:.*]]: vector<4xf32>95func.func @transfer_to_maskedload_addrspace(%mem : memref<8x8xf32, #gpu.address_space<workgroup>>, %idx : index, %mask : vector<4xi1>, %passthru : vector<4xf32>) -> vector<4xf32> {96 %res = vector.maskedload %mem[%idx, %idx], %mask, %passthru : memref<8x8xf32, #gpu.address_space<workgroup>>, vector<4xi1>, vector<4xf32> into vector<4xf32>97 return %res : vector<4xf32>98}99// CHECK: %[[RES:.*]] = vector.maskedload %[[ARG0]][%[[ARG1]], %[[ARG1]]], %[[ARG2]], %[[ARG3]]100// CHECK: return %[[RES]] : vector<4xf32>101 102// -----103 104// CHECK-LABEL: func @transfer_scalar(105// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>106// CHECK-SAME: %[[ARG1:.*]]: index107// CHECK-SAME: %[[ARG2:.*]]: vector<1xi1>108// CHECK-SAME: %[[ARG3:.*]]: vector<1xf32>109func.func @transfer_scalar(%mem : memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>, %idx : index, %mask : vector<1xi1>, %passthru : vector<1xf32>) -> vector<1xf32> {110 %res = vector.maskedload %mem[%idx, %idx], %mask, %passthru111 : memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>, vector<1xi1>, vector<1xf32> into vector<1xf32>112 return %res : vector<1xf32>113}114// CHECK: %[[IF:.*]] = scf.if115// CHECK: %[[LOAD:.*]] = vector.load %[[ARG0]][%[[ARG1]], %[[ARG1]]]116// CHECK: %[[SELECT:.*]] = arith.select %arg2, %[[LOAD]], %[[ARG3]]117 118// -----119 120func.func @transfer_to_maskedload_fatrawbuffer(%mem : memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>, %idx : index, %mask : vector<4xi1>, %passthru : vector<4xf32>) -> vector<4xf32> {121 %res = vector.maskedload %mem[%idx, %idx], %mask, %passthru : memref<8x8xf32, #amdgpu.address_space<fat_raw_buffer>>, vector<4xi1>, vector<4xf32> into vector<4xf32>122 return %res : vector<4xf32>123}124 125// -----126 127// CHECK-LABEL: func.func @full_select_maskedload_fatrawbuffer_to_load128func.func @full_select_maskedload_fatrawbuffer_to_load(%arg0: memref<8x8xf16, #amdgpu.address_space<fat_raw_buffer>>, %arg1: index, %arg2: i1, %arg3: vector<4xf16>) -> vector<4xf16> {129 %0 = vector.broadcast %arg2 : i1 to vector<4xi1>130 %1 = vector.maskedload %arg0[%arg1, %arg1], %0, %arg3 : memref<8x8xf16, #amdgpu.address_space<fat_raw_buffer>>, vector<4xi1>, vector<4xf16> into vector<4xf16>131 return %1 : vector<4xf16>132}133// CHECK-NOT: vector.maskedload134// CHECK: vector.load135// CHECK: arith.select136 137// -----138 139// CHECK-LABEL: func.func @full_select_maskedload_to_load140// CHECK-SAME: %[[MEM:.+]]: memref<8x8xf16>,141// CHECK-SAME: %[[IDX:.+]]: index,142// CHECK-SAME: %[[PRED:.+]]: i1,143// CHECK-SAME: %[[PASSTHRU:.+]]: vector<4xf16>)144func.func @full_select_maskedload_to_load(%arg0: memref<8x8xf16>, %arg1: index, %arg2: i1, %arg3: vector<4xf16>) -> vector<4xf16> {145 %0 = vector.broadcast %arg2 : i1 to vector<4xi1>146 %1 = vector.maskedload %arg0[%arg1, %arg1], %0, %arg3 : memref<8x8xf16>, vector<4xi1>, vector<4xf16> into vector<4xf16>147 return %1 : vector<4xf16>148}149// CHECK-NOT: vector.maskedload150// CHECK: scf.if %[[PRED]]151// CHECK: %[[LOAD:.+]] = vector.load152// CHECK: scf.yield %[[LOAD]]153// CHECK: else154// CHECK: scf.yield %[[PASSTHRU]]155 156// -----157 158// CHECK-LABEL: func.func @full_mask_maskedstore_to_store159// CHECK-SAME: %[[MEM:.+]]: memref<8x8xf16>,160// CHECK-SAME: %[[IDX:.+]]: index,161// CHECK-SAME: %[[PRED:.+]]: i1,162func.func @full_mask_maskedstore_to_store(%arg0: memref<8x8xf16>, %arg1: index, %arg2: i1, %arg3: vector<4xf16>) {163 %0 = vector.broadcast %arg2 : i1 to vector<4xi1>164 vector.maskedstore %arg0[%arg1, %arg1], %0, %arg3 : memref<8x8xf16>, vector<4xi1>, vector<4xf16>165 return166}167// CHECK-NOT: vector.maskedstore168// CHECK: scf.if %[[PRED]]169// CHECK: vector.store170