670 lines · plain
1// RUN: mlir-opt %s -arm-sme-vector-legalization -cse -canonicalize -split-input-file | FileCheck %s2 3// CHECK-LABEL: @outerproduct_f32_scalable_8x8_no_acc(4// CHECK-SAME: %[[LHS:.*]]: vector<[8]xf32>,5// CHECK-SAME: %[[RHS:.*]]: vector<[8]xf32>)6// CHECK-SAME: -> (vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>)7func.func @outerproduct_f32_scalable_8x8_no_acc(%lhs: vector<[8]xf32>, %rhs: vector<[8]xf32>) -> vector<[8]x[8]xf32>8{9 // CHECK-DAG: %[[LHS_0:.*]] = vector.scalable.extract %[[LHS]][0] : vector<[4]xf32> from vector<[8]xf32>10 // CHECK-DAG: %[[RHS_0:.*]] = vector.scalable.extract %[[RHS]][0] : vector<[4]xf32> from vector<[8]xf32>11 // CHECK-DAG: %[[LHS_1:.*]] = vector.scalable.extract %[[LHS]][4] : vector<[4]xf32> from vector<[8]xf32>12 // CHECK-DAG: %[[RHS_1:.*]] = vector.scalable.extract %[[RHS]][4] : vector<[4]xf32> from vector<[8]xf32>13 // CHECK-DAG: %[[TOP_LEFT:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_0]] : vector<[4]xf32>, vector<[4]xf32>14 // CHECK-DAG: %[[TOP_RIGHT:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_1]] : vector<[4]xf32>, vector<[4]xf32>15 // CHECK-DAG: %[[BOTTOM_LEFT:.*]] = vector.outerproduct %[[LHS_1]], %[[RHS_0]] : vector<[4]xf32>, vector<[4]xf32>16 // CHECK-DAG: %[[BOTTOM_RIGHT:.*]] = vector.outerproduct %[[LHS_1]], %[[RHS_1]] : vector<[4]xf32>, vector<[4]xf32>17 // CHECK-NEXT: return %[[TOP_LEFT]], %[[TOP_RIGHT]], %[[BOTTOM_LEFT]], %[[BOTTOM_RIGHT]] : vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>18 %0 = vector.outerproduct %lhs, %rhs : vector<[8]xf32>, vector<[8]xf32>19 return %0 : vector<[8]x[8]xf32>20}21 22// -----23 24// CHECK-LABEL: @outerproduct_f32_scalable_4x16_acc(25// CHECK-SAME: %[[LHS:.*]]: vector<[4]xf32>,26// CHECK-SAME: %[[RHS:.*]]: vector<[16]xf32>,27// CHECK-SAME: %[[ACC_0:[A-Za-z0-9]*]]: vector<[4]x[4]xf32>,28// CHECK-SAME: %[[ACC_1:[A-Za-z0-9]*]]: vector<[4]x[4]xf32>,29// CHECK-SAME: %[[ACC_2:[A-Za-z0-9]*]]: vector<[4]x[4]xf32>,30// CHECK-SAME: %[[ACC_3:[A-Za-z0-9]*]]: vector<[4]x[4]xf32>)31// CHECK-SAME: -> (vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>)32func.func @outerproduct_f32_scalable_4x16_acc(%lhs: vector<[4]xf32>, %rhs: vector<[16]xf32>, %acc: vector<[4]x[16]xf32>) -> vector<[4]x[16]xf32>33{34 // CHECK-DAG: %[[LHS_0:.*]] = vector.scalable.extract %[[LHS]][0] : vector<[4]xf32> from vector<[4]xf32>35 // CHECK-DAG: %[[RHS_0:.*]] = vector.scalable.extract %[[RHS]][0] : vector<[4]xf32> from vector<[16]xf32>36 // CHECK-DAG: %[[RHS_1:.*]] = vector.scalable.extract %[[RHS]][4] : vector<[4]xf32> from vector<[16]xf32>37 // CHECK-DAG: %[[RHS_2:.*]] = vector.scalable.extract %[[RHS]][8] : vector<[4]xf32> from vector<[16]xf32>38 // CHECK-DAG: %[[RHS_3:.*]] = vector.scalable.extract %[[RHS]][12] : vector<[4]xf32> from vector<[16]xf32>39 // CHECK-DAG: %[[RES_0:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_0]], %[[ACC_0]] {kind = #vector.kind<add>} : vector<[4]xf32>, vector<[4]xf32>40 // CHECK-DAG: %[[RES_1:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_1]], %[[ACC_1]] {kind = #vector.kind<add>} : vector<[4]xf32>, vector<[4]xf32>41 // CHECK-DAG: %[[RES_2:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_2]], %[[ACC_2]] {kind = #vector.kind<add>} : vector<[4]xf32>, vector<[4]xf32>42 // CHECK-DAG: %[[RES_3:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_3]], %[[ACC_3]] {kind = #vector.kind<add>} : vector<[4]xf32>, vector<[4]xf32>43 // CHECK-NEXT: return %[[RES_0]], %[[RES_1]], %[[RES_2]], %[[RES_3]] : vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>44 %0 = vector.outerproduct %lhs, %rhs, %acc : vector<[4]xf32>, vector<[16]xf32>45 return %0 : vector<[4]x[16]xf32>46}47 48// -----49 50// CHECK-LABEL: @outerproduct_f32_masked_scalable_16x4(51// CHECK-SAME: %[[LHS:.*]]: vector<[16]xf32>,52// CHECK-SAME: %[[RHS:.*]]: vector<[4]xf32>,53// CHECK-SAME: %[[LHS_DIM:.*]]: index,54// CHECK-SAME: %[[RHS_DIM:.*]]: index)55// CHECK-SAME: -> (vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>)56func.func @outerproduct_f32_masked_scalable_16x4(%lhs: vector<[16]xf32>, %rhs: vector<[4]xf32>, %lhs_dim: index, %rhs_dim: index) -> vector<[16]x[4]xf32>57{58 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale59 // CHECK-DAG: %[[MINUS_4:.*]] = arith.constant -4 : index60 // CHECK-DAG: %[[MINUS_8:.*]] = arith.constant -8 : index61 // CHECK-DAG: %[[MINUS_12:.*]] = arith.constant -12 : index62 // CHECK-DAG: %[[MINUS_4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[MINUS_4]] : index63 // CHECK-DAG: %[[MINUS_8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[MINUS_8]] : index64 // CHECK-DAG: %[[MINUS_12_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[MINUS_12]] : index65 // CHECK-DAG: %[[LHS_0:.*]] = vector.scalable.extract %[[LHS]][0] : vector<[4]xf32> from vector<[16]xf32>66 // CHECK-DAG: %[[LHS_1:.*]] = vector.scalable.extract %[[LHS]][4] : vector<[4]xf32> from vector<[16]xf32>67 // CHECK-DAG: %[[LHS_2:.*]] = vector.scalable.extract %[[LHS]][8] : vector<[4]xf32> from vector<[16]xf32>68 // CHECK-DAG: %[[LHS_3:.*]] = vector.scalable.extract %[[LHS]][12] : vector<[4]xf32> from vector<[16]xf32>69 // CHECK-DAG: %[[RHS_0:.*]] = vector.scalable.extract %[[RHS]][0] : vector<[4]xf32> from vector<[4]xf32>70 // CHECK-DAG: %[[MASK_0:.*]] = vector.create_mask %[[LHS_DIM]], %[[RHS_DIM]] : vector<[4]x[4]xi1>71 // CHECK-DAG: %[[TILE_1_LHS_DIM:.*]] = arith.addi %[[LHS_DIM]], %[[MINUS_4_VSCALE]] : index72 // CHECK-DAG: %[[MASK_1:.*]] = vector.create_mask %[[TILE_1_LHS_DIM]], %[[RHS_DIM]] : vector<[4]x[4]xi1>73 // CHECK-DAG: %[[TILE_2_LHS_DIM:.*]] = arith.addi %[[LHS_DIM]], %[[MINUS_8_VSCALE]] : index74 // CHECK-DAG: %[[MASK_2:.*]] = vector.create_mask %[[TILE_2_LHS_DIM]], %[[RHS_DIM]] : vector<[4]x[4]xi1>75 // CHECK-DAG: %[[TILE_3_LHS_DIM:.*]] = arith.addi %[[LHS_DIM]], %[[MINUS_12_VSCALE]] : index76 // CHECK-DAG: %[[MASK_3:.*]] = vector.create_mask %[[TILE_3_LHS_DIM]], %[[RHS_DIM]] : vector<[4]x[4]xi1>77 // CHECK-DAG: %[[RES_0:.*]] = vector.mask %[[MASK_0]] { vector.outerproduct %[[LHS_0]], %[[RHS_0]] : vector<[4]xf32>, vector<[4]xf32> } : vector<[4]x[4]xi1> -> vector<[4]x[4]xf32>78 // CHECK-DAG: %[[RES_1:.*]] = vector.mask %[[MASK_1]] { vector.outerproduct %[[LHS_1]], %[[RHS_0]] : vector<[4]xf32>, vector<[4]xf32> } : vector<[4]x[4]xi1> -> vector<[4]x[4]xf32>79 // CHECK-DAG: %[[RES_2:.*]] = vector.mask %[[MASK_2]] { vector.outerproduct %[[LHS_2]], %[[RHS_0]] : vector<[4]xf32>, vector<[4]xf32> } : vector<[4]x[4]xi1> -> vector<[4]x[4]xf32>80 // CHECK-DAG: %[[RES_3:.*]] = vector.mask %[[MASK_3]] { vector.outerproduct %[[LHS_3]], %[[RHS_0]] : vector<[4]xf32>, vector<[4]xf32> } : vector<[4]x[4]xi1> -> vector<[4]x[4]xf32>81 // CHECK-NEXT: return %[[RES_0]], %[[RES_1]], %[[RES_2]], %[[RES_3]] : vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>82 %mask = vector.create_mask %lhs_dim, %rhs_dim : vector<[16]x[4]xi1>83 %0 = vector.mask %mask { vector.outerproduct %lhs, %rhs : vector<[16]xf32>, vector<[4]xf32> } : vector<[16]x[4]xi1> -> vector<[16]x[4]xf32>84 return %0 : vector<[16]x[4]xf32>85}86 87// -----88 89/// This demonstrates a rectangular tiling that uses all f64 accumulators.90 91// CHECK-LABEL: @outerproduct_f64_scalable_8x4_no_acc(92// CHECK-SAME: %[[LHS:.*]]: vector<[8]xf64>,93// CHECK-SAME: %[[RHS:.*]]: vector<[4]xf64>)94// CHECK-SAME: -> (vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>)95func.func @outerproduct_f64_scalable_8x4_no_acc(%lhs: vector<[8]xf64>, %rhs: vector<[4]xf64>) -> vector<[8]x[4]xf64>96{97 // CHECK-DAG: %[[LHS_0:.*]] = vector.scalable.extract %[[LHS]][0] : vector<[2]xf64> from vector<[8]xf64>98 // CHECK-DAG: %[[LHS_1:.*]] = vector.scalable.extract %[[LHS]][2] : vector<[2]xf64> from vector<[8]xf64>99 // CHECK-DAG: %[[LHS_2:.*]] = vector.scalable.extract %[[LHS]][4] : vector<[2]xf64> from vector<[8]xf64>100 // CHECK-DAG: %[[LHS_3:.*]] = vector.scalable.extract %[[LHS]][6] : vector<[2]xf64> from vector<[8]xf64>101 // CHECK-DAG: %[[RHS_0:.*]] = vector.scalable.extract %[[RHS]][0] : vector<[2]xf64> from vector<[4]xf64>102 // CHECK-DAG: %[[RHS_1:.*]] = vector.scalable.extract %[[RHS]][2] : vector<[2]xf64> from vector<[4]xf64>103 // CHECK-DAG: %[[RES_0:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_0]] : vector<[2]xf64>, vector<[2]xf64>104 // CHECK-DAG: %[[RES_1:.*]] = vector.outerproduct %[[LHS_0]], %[[RHS_1]] : vector<[2]xf64>, vector<[2]xf64>105 // CHECK-DAG: %[[RES_2:.*]] = vector.outerproduct %[[LHS_1]], %[[RHS_0]] : vector<[2]xf64>, vector<[2]xf64>106 // CHECK-DAG: %[[RES_3:.*]] = vector.outerproduct %[[LHS_1]], %[[RHS_1]] : vector<[2]xf64>, vector<[2]xf64>107 // CHECK-DAG: %[[RES_4:.*]] = vector.outerproduct %[[LHS_2]], %[[RHS_0]] : vector<[2]xf64>, vector<[2]xf64>108 // CHECK-DAG: %[[RES_5:.*]] = vector.outerproduct %[[LHS_2]], %[[RHS_1]] : vector<[2]xf64>, vector<[2]xf64>109 // CHECK-DAG: %[[RES_6:.*]] = vector.outerproduct %[[LHS_3]], %[[RHS_0]] : vector<[2]xf64>, vector<[2]xf64>110 // CHECK-DAG: %[[RES_7:.*]] = vector.outerproduct %[[LHS_3]], %[[RHS_1]] : vector<[2]xf64>, vector<[2]xf64>111 // CHECK-NEXT: return %[[RES_0]], %[[RES_1]], %[[RES_2]], %[[RES_3]], %[[RES_4]], %[[RES_5]], %[[RES_6]], %[[RES_7]] : vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>, vector<[2]x[2]xf64>112 %0 = vector.outerproduct %lhs, %rhs : vector<[8]xf64>, vector<[4]xf64>113 return %0 : vector<[8]x[4]xf64>114}115 116// -----117 118// CHECK-LABEL: @transfer_read_f32_scalable_8x8(119// CHECK-SAME: %[[SRC:.*]]: memref<?x?xi32>)120// CHECK-SAME: -> (vector<[4]x[4]xi32>, vector<[4]x[4]xi32>, vector<[4]x[4]xi32>, vector<[4]x[4]xi32>)121func.func @transfer_read_f32_scalable_8x8(%src: memref<?x?xi32>) -> vector<[8]x[8]xi32>122{123 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index124 // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index125 // CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32126 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale127 // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index128 // CHECK-DAG: %[[TOP_LEFT:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C0]]], %[[C0_I32]] {in_bounds = [true, true]} : memref<?x?xi32>, vector<[4]x[4]xi32>129 // CHECK-DAG: %[[TOP_RIGHT:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C4_VSCALE]]], %[[C0_I32]] {in_bounds = [true, true]} : memref<?x?xi32>, vector<[4]x[4]xi32>130 // CHECK-DAG: %[[BOTTOM_LEFT:.*]] = vector.transfer_read %[[SRC]][%[[C4_VSCALE]], %[[C0]]], %[[C0_I32]] {in_bounds = [true, true]} : memref<?x?xi32>, vector<[4]x[4]xi32>131 // CHECK-DAG: %[[BOTTOM_RIGHT:.*]] = vector.transfer_read %[[SRC]][%[[C4_VSCALE]], %[[C4_VSCALE]]], %[[C0_I32]] {in_bounds = [true, true]} : memref<?x?xi32>, vector<[4]x[4]xi32>132 // CHECK-NEXT: return %[[TOP_LEFT]], %[[TOP_RIGHT]], %[[BOTTOM_LEFT]], %[[BOTTOM_RIGHT]] : vector<[4]x[4]xi32>, vector<[4]x[4]xi32>, vector<[4]x[4]xi32>, vector<[4]x[4]xi32>133 %c0 = arith.constant 0 : index134 %pad = arith.constant 0 : i32135 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xi32>, vector<[8]x[8]xi32>136 return %0 : vector<[8]x[8]xi32>137}138 139// -----140 141// CHECK-LABEL: @transfer_read_i16_scalable_8x16_masked(142// CHECK-SAME: %[[SRC:.*]]: memref<?x?xi16>,143// CHECK-SAME: %[[DIM0:.*]]: index,144// CHECK-SAME: %[[DIM1:.*]]: index)145// CHECK-SAME: -> (vector<[8]x[8]xi16>, vector<[8]x[8]xi16>)146func.func @transfer_read_i16_scalable_8x16_masked(%src: memref<?x?xi16>, %dim0: index, %dim1: index) -> vector<[8]x[16]xi16>147{148 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index149 // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index150 // CHECK-DAG: %[[MINUS_8:.*]] = arith.constant -8 : index151 // CHECK-DAG: %[[C0_I16:.*]] = arith.constant 0 : i16152 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale153 // CHECK-DAG: %[[MINUS_8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[MINUS_8]] : index154 // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index155 // CHECK-DAG: %[[RIGHT_DIM_1:.*]] = arith.addi %[[DIM1]], %[[MINUS_8_VSCALE]] : index156 // CHECK-DAG: %[[LEFT_MASK:.*]] = vector.create_mask %[[DIM0]], %[[DIM1]] : vector<[8]x[8]xi1>157 // CHECK-DAG: %[[RIGHT_MASK:.*]] = vector.create_mask %[[DIM0]], %[[RIGHT_DIM_1]] : vector<[8]x[8]xi1>158 // CHECK-DAG: %[[LEFT:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C0]]], %[[C0_I16]], %[[LEFT_MASK]] {in_bounds = [true, true]} : memref<?x?xi16>, vector<[8]x[8]xi16>159 // CHECK-DAG: %[[RIGHT:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C8_VSCALE]]], %[[C0_I16]], %[[RIGHT_MASK]] {in_bounds = [true, true]} : memref<?x?xi16>, vector<[8]x[8]xi16>160 // CHECK-NEXT: return %[[LEFT]], %[[RIGHT]] : vector<[8]x[8]xi16>, vector<[8]x[8]xi16>161 %c0 = arith.constant 0 : index162 %pad = arith.constant 0 : i16163 %mask = vector.create_mask %dim0, %dim1 : vector<[8]x[16]xi1>164 %0 = vector.transfer_read %src[%c0, %c0], %pad, %mask {in_bounds = [true, true]} : memref<?x?xi16>, vector<[8]x[16]xi16>165 return %0 : vector<[8]x[16]xi16>166}167 168// -----169 170// CHECK-LABEL: @transfer_write_f16_scalable_16x8(171// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf16>,172// CHECK-SAME: %[[TOP:.*]]: vector<[8]x[8]xf16>,173// CHECK-SAME: %[[BOTTOM:.*]]: vector<[8]x[8]xf16>)174func.func @transfer_write_f16_scalable_16x8(%dest: memref<?x?xf16>, %vec: vector<[16]x[8]xf16>)175{176 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index177 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index178 // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index179 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale180 // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index181 // CHECK-NEXT: scf.for %[[I:.*]] = %[[C0]] to %[[C8_VSCALE]] step %[[C1]] {182 // CHECK-NEXT: %[[TOP_SLICE:.*]] = vector.extract %[[TOP]][%[[I]]] : vector<[8]xf16> from vector<[8]x[8]xf16>183 // CHECK-NEXT: vector.transfer_write %[[TOP_SLICE]], %[[DEST]][%[[I]], %[[C0]]] {in_bounds = [true]} : vector<[8]xf16>, memref<?x?xf16>184 // CHECK-NEXT: %[[BOTTOM_I:.*]] = arith.addi %[[C8_VSCALE]], %[[I]] : index185 // CHECK-NEXT: %[[BOTTOM_SLICE:.*]] = vector.extract %[[BOTTOM]][%[[I]]] : vector<[8]xf16> from vector<[8]x[8]xf16>186 // CHECK-NEXT: vector.transfer_write %[[BOTTOM_SLICE]], %[[DEST]][%[[BOTTOM_I]], %[[C0]]] {in_bounds = [true]} : vector<[8]xf16>, memref<?x?xf16>187 // CHECK-NEXT: }188 // CHECK-NEXT: return189 %c0 = arith.constant 0 : index190 vector.transfer_write %vec, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[8]xf16>, memref<?x?xf16>191 return192}193 194// -----195 196/// This is already a legal type. It should be ignored.197 198// CHECK-LABEL: @transfer_write_i8_scalable_16x16_masked199func.func @transfer_write_i8_scalable_16x16_masked(%dest: memref<?x?xi8>, %vec: vector<[16]x[16]xi8>, %dim0: index, %dim1: index)200{201 // CHECK: vector.transfer_write {{.*}} : vector<[16]x[16]xi8>, memref<?x?xi8>202 %c0 = arith.constant 0 : index203 %mask = vector.create_mask %dim0, %dim0 : vector<[16]x[16]xi1>204 vector.transfer_write %vec, %dest[%c0, %c0], %mask {in_bounds = [true, true]} : vector<[16]x[16]xi8>, memref<?x?xi8>205 return206}207 208// -----209 210// CHECK-LABEL: @transfer_write_f32_scalable_8x8_masked(211// CHECK-SAME: %[[DEST:[a-z0-9]+]]: memref<?x?xf32>,212// CHECK-SAME: %[[DIM_0:[a-z0-9]+]]: index,213// CHECK-SAME: %[[DIM_1:[a-z0-9]+]]: index,214// CHECK-SAME: %[[TILE_0:[a-z0-9]+]]: vector<[4]x[4]xf32>,215// CHECK-SAME: %[[TILE_1:[a-z0-9]+]]: vector<[4]x[4]xf32>,216// CHECK-SAME: %[[TILE_2:[a-z0-9]+]]: vector<[4]x[4]xf32>,217// CHECK-SAME: %[[TILE_3:[a-z0-9]+]]: vector<[4]x[4]xf32>)218func.func @transfer_write_f32_scalable_8x8_masked(%dest: memref<?x?xf32>, %dim0: index, %dim1: index, %vec: vector<[8]x[8]xf32>)219{220 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index221 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index222 // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index223 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale224 // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index225 // CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[DIM_0]], %[[DIM_1]] : vector<[8]x[8]xi1>226 // CHECK-NEXT: scf.for %[[I:.*]] = %[[C0]] to %[[C4_VSCALE]] step %[[C1]] {227 // CHECK-NEXT: %[[UPPER_SLICE_MASK:.*]] = vector.extract %[[MASK]][%[[I]]] : vector<[8]xi1> from vector<[8]x[8]xi1>228 // CHECK-NEXT: %[[TILE_0_SLICE_MASK:.*]] = vector.scalable.extract %[[UPPER_SLICE_MASK]][0] : vector<[4]xi1> from vector<[8]xi1>229 // CHECK-NEXT: %[[TILE_0_SLICE:.*]] = vector.extract %[[TILE_0]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>230 // CHECK-NEXT: vector.transfer_write %[[TILE_0_SLICE]], %[[DEST]][%[[I]], %[[C0]]], %[[TILE_0_SLICE_MASK]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>231 // CHECK-NEXT: %[[TILE_1_SLICE_MASK:.*]] = vector.scalable.extract %[[UPPER_SLICE_MASK]][4] : vector<[4]xi1> from vector<[8]xi1>232 // CHECK-NEXT: %[[TILE_1_SLICE:.*]] = vector.extract %[[TILE_1]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>233 // CHECK-NEXT: vector.transfer_write %[[TILE_1_SLICE]], %[[DEST]][%[[I]], %[[C4_VSCALE]]], %[[TILE_1_SLICE_MASK]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>234 // CHECK-NEXT: %[[LOWER_SLICE_I:.*]] = arith.addi %[[C4_VSCALE]], %[[I]] : index235 // CHECK-NEXT: %[[LOWER_SLICE_MASK:.*]] = vector.extract %[[MASK]][%[[LOWER_SLICE_I]]] : vector<[8]xi1> from vector<[8]x[8]xi1>236 // CHECK-NEXT: %[[TILE_2_SLICE_MASK:.*]] = vector.scalable.extract %[[LOWER_SLICE_MASK]][0] : vector<[4]xi1> from vector<[8]xi1>237 // CHECK-NEXT: %[[TILE_2_SLICE:.*]] = vector.extract %[[TILE_2]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>238 // CHECK-NEXT: vector.transfer_write %[[TILE_2_SLICE]], %[[DEST]][%[[LOWER_SLICE_I]], %[[C0]]], %[[TILE_2_SLICE_MASK]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>239 // CHECK-NEXT: %[[TILE_3_SLICE_MASK:.*]] = vector.scalable.extract %[[LOWER_SLICE_MASK]][4] : vector<[4]xi1> from vector<[8]xi1>240 // CHECK-NEXT: %[[TILE_3_SLICE:.*]] = vector.extract %[[TILE_3]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>241 // CHECK-NEXT: vector.transfer_write %[[TILE_3_SLICE:.*]], %[[DEST]][%[[LOWER_SLICE_I]], %[[C4_VSCALE]]], %[[TILE_3_SLICE_MASK]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>242 // CHECK-NEXT: }243 %c0 = arith.constant 0 : index244 %mask = vector.create_mask %dim0, %dim1 : vector<[8]x[8]xi1>245 vector.transfer_write %vec, %dest[%c0, %c0], %mask {in_bounds = [true, true]} : vector<[8]x[8]xf32>, memref<?x?xf32>246 return247}248 249// -----250 251// Tensor semantics are not supported for the store loop lowering.252 253// CHECK-LABEL: @negative_transfer_write_f32_scalable_8x8_tensor254// CHECK-NOT: scf.for255func.func @negative_transfer_write_f32_scalable_8x8_tensor(%dest: tensor<?x?xf32>, %vec: vector<[8]x[8]xf32>)256{257 %c0 = arith.constant 0 : index258 vector.transfer_write %vec, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[8]x[8]xf32>, tensor<?x?xf32>259 return260}261 262// -----263 264#transpose = affine_map<(d0, d1) -> (d1, d0)>265 266// Transposes are not supported for the store loop lowering.267 268// CHECK-LABEL: @negative_transfer_write_f32_scalable_8x8_tensor269// CHECK-NOT: scf.for270func.func @negative_transfer_write_f32_scalable_8x8_tensor(%dest: tensor<?x?xf32>, %dim0: index, %dim1: index, %vec: vector<[8]x[8]xf32>)271{272 %c0 = arith.constant 0 : index273 %mask = vector.create_mask %dim0, %dim1 : vector<[8]x[8]xi1>274 vector.transfer_write %vec, %dest[%c0, %c0], %mask {permutation_map = #transpose, in_bounds = [true, true]} : vector<[8]x[8]xf32>, tensor<?x?xf32>275 return276}277 278// -----279 280// Masked writes where any dimension of the mask is > 16 are not supported for the store loop lowering.281 282// CHECK-LABEL: @negative_transfer_write_f32_scalable_32x32283// CHECK-NOT: scf.for284func.func @negative_transfer_write_f32_scalable_32x32(%dest: memref<?x?xf32>, %dim0: index, %dim1: index, %vec: vector<[32]x[32]xf32>)285{286 %c0 = arith.constant 0 : index287 %mask = vector.create_mask %dim0, %dim1 : vector<[32]x[32]xi1>288 vector.transfer_write %vec, %dest[%c0, %c0], %mask {in_bounds = [true, true]} : vector<[32]x[32]xf32>, memref<?x?xf32>289 return290}291 292// -----293 294#transpose = affine_map<(d0, d1) -> (d1, d0)>295 296// CHECK-LABEL: @transpose_f32_scalable_4x16_via_read(297// CHECK-SAME: %[[SRC:.*]]: memref<?x?xf32>,298// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>)299func.func @transpose_f32_scalable_4x16_via_read(%src: memref<?x?xf32>, %dest: memref<?x?xf32>)300{301 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index302 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index303 // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index304 // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index305 // CHECK-DAG: %[[C12:.*]] = arith.constant 12 : index306 // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32307 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale308 // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index309 // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index310 // CHECK-DAG: %[[C12_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C12]] : index311 // CHECK-DAG: %[[TILE_0:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<[4]x[4]xf32>312 // CHECK-DAG: %[[TILE_1:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C4_VSCALE]]], %[[PAD]] {in_bounds = [true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<[4]x[4]xf32>313 // CHECK-DAG: %[[TILE_2:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C8_VSCALE]]], %[[PAD]] {in_bounds = [true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<[4]x[4]xf32>314 // CHECK-DAG: %[[TILE_3:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C12_VSCALE]]], %[[PAD]] {in_bounds = [true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<[4]x[4]xf32>315 // CHECK-NEXT: scf.for %[[I:.*]] = %[[C0]] to %[[C4_VSCALE]] step %[[C1]] {316 // CHECK-NEXT: %[[TILE_0_SLICE:.*]] = vector.extract %[[TILE_0]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>317 // CHECK-NEXT: vector.transfer_write %[[TILE_0_SLICE]], %[[DEST]][%[[I]], %[[C0]]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>318 // CHECK-NEXT: %[[TILE_1_I:.*]] = arith.addi %[[C4_VSCALE]], %[[I]] : index319 // CHECK-NEXT: %[[TILE_1_SLICE:.*]] = vector.extract %[[TILE_1]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>320 // CHECK-NEXT: vector.transfer_write %[[TILE_1_SLICE]], %[[DEST]][%[[TILE_1_I]], %[[C0]]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>321 // CHECK-NEXT: %[[TILE_2_I:.*]] = arith.addi %[[C8_VSCALE]], %[[I]] : index322 // CHECK-NEXT: %[[TILE_2_SLICE:.*]] = vector.extract %[[TILE_2]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>323 // CHECK-NEXT: vector.transfer_write %[[TILE_2_SLICE]], %[[DEST]][%[[TILE_2_I]], %[[C0]]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>324 // CHECK-NEXT: %[[TILE_3_I:.*]] = arith.addi %[[C12_VSCALE]], %[[I]] : index325 // CHECK-NEXT: %[[TILE_3_SLICE:.*]] = vector.extract %[[TILE_3]][%[[I]]] : vector<[4]xf32> from vector<[4]x[4]xf32>326 // CHECK-NEXT: vector.transfer_write %[[TILE_3_SLICE]], %[[DEST]][%[[TILE_3_I]], %[[C0]]] {in_bounds = [true]} : vector<[4]xf32>, memref<?x?xf32>327 // CHECK-NEXT: }328 // CHECK-NEXT: return329 %c0 = arith.constant 0 : index330 %pad = arith.constant 0.0 : f32331 %0 = vector.transfer_read %src[%c0, %c0], %pad {permutation_map = #transpose, in_bounds = [true, true]} : memref<?x?xf32>, vector<[16]x[4]xf32>332 vector.transfer_write %0, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[4]xf32>, memref<?x?xf32>333 return334}335 336// -----337 338#transpose = affine_map<(d0, d1) -> (d1, d0)>339 340// CHECK-LABEL: @transpose_f32_scalable_4x16_via_write(341// CHECK-SAME: %[[SRC:.*]]: memref<?x?xf32>,342// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>)343func.func @transpose_f32_scalable_4x16_via_write(%src: memref<?x?xf32>, %dest: memref<?x?xf32>)344{345 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index346 // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index347 // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index348 // CHECK-DAG: %[[C12:.*]] = arith.constant 12 : index349 // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32350 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale351 // CHECK-DAG: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index352 // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index353 // CHECK-DAG: %[[C12_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C12]] : index354 // CHECK-DAG: %[[TILE_0:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>355 // CHECK-DAG: %[[TILE_1:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C4_VSCALE]]], %[[PAD]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>356 // CHECK-DAG: %[[TILE_2:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C8_VSCALE]]], %[[PAD]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>357 // CHECK-DAG: %[[TILE_3:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C12_VSCALE]]], %[[PAD]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>358 // CHECK-DAG: vector.transfer_write %[[TILE_0]], %[[DEST]][%[[C0]], %[[C0]]] {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<[4]x[4]xf32>, memref<?x?xf32>359 // CHECK-DAG: vector.transfer_write %[[TILE_1]], %[[DEST]][%[[C4_VSCALE]], %[[C0]]] {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<[4]x[4]xf32>, memref<?x?xf32>360 // CHECK-DAG: vector.transfer_write %[[TILE_2]], %[[DEST]][%[[C8_VSCALE]], %[[C0]]] {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<[4]x[4]xf32>, memref<?x?xf32>361 // CHECK-DAG: vector.transfer_write %[[TILE_3]], %[[DEST]][%[[C12_VSCALE]], %[[C0]]] {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<[4]x[4]xf32>, memref<?x?xf32>362 // CHECK-NEXT: return363 %c0 = arith.constant 0 : index364 %pad = arith.constant 0.0 : f32365 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[16]xf32>366 vector.transfer_write %0, %dest[%c0, %c0] {permutation_map = #transpose, in_bounds = [true, true]} : vector<[4]x[16]xf32>, memref<?x?xf32>367 return368}369 370// -----371 372// CHECK-LABEL: @extract_from_vector_create_mask_non_constant_dim(373// CHECK-SAME: %[[DIM0:[a-z0-9]+]]: index,374// CHECK-SAME: %[[DIM1:[a-z0-9]+]]: index,375// CHECK-SAME: %[[DIM2:[a-z0-9]+]]: index)376func.func @extract_from_vector_create_mask_non_constant_dim(%dim0: index, %dim1: index, %dim2: index) -> vector<[4]x[4]xi1> {377 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index378 // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index379 // CHECK-NEXT: %[[DIM0_CMP:.*]] = arith.cmpi sgt, %[[DIM0]], %[[C2]] : index380 // CHECK-NEXT: %[[NEW_DIM0:.*]] = arith.select %[[DIM0_CMP]], %[[DIM1]], %[[C0]] : index381 // CHECK-NEXT: %[[EXTRACT:.*]] = vector.create_mask %[[NEW_DIM0]], %[[DIM2]] : vector<[4]x[4]xi1>382 // CHECK-NEXT: return %[[EXTRACT]]383 %mask = vector.create_mask %dim0, %dim1, %dim2 : vector<4x[4]x[4]xi1>384 %extract = vector.extract %mask[2] : vector<[4]x[4]xi1> from vector<4x[4]x[4]xi1>385 return %extract : vector<[4]x[4]xi1>386}387 388// -----389 390// CHECK-LABEL: @non_constant_extract_from_vector_create_mask_non_constant(391// CHECK-SAME: %[[INDEX:[a-z0-9]+]]: index,392// CHECK-SAME: %[[DIM0:[a-z0-9]+]]: index,393// CHECK-SAME: %[[DIM1:[a-z0-9]+]]: index,394// CHECK-SAME: %[[DIM2:[a-z0-9]+]]: index)395func.func @non_constant_extract_from_vector_create_mask_non_constant(%index: index, %dim0: index, %dim1: index, %dim2: index) -> vector<[4]x[4]xi1> {396 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index397 // CHECK-NEXT: %[[DIM0_CMP:.*]] = arith.cmpi slt, %[[INDEX]], %[[DIM0]] : index398 // CHECK-NEXT: %[[NEW_DIM0:.*]] = arith.select %[[DIM0_CMP]], %[[DIM1]], %[[C0]] : index399 // CHECK-NEXT: %[[EXTRACT:.*]] = vector.create_mask %[[NEW_DIM0]], %[[DIM2]] : vector<[4]x[4]xi1>400 // CHECK-NEXT: return %[[EXTRACT]]401 %mask = vector.create_mask %dim0, %dim1, %dim2 : vector<4x[4]x[4]xi1>402 %extract = vector.extract %mask[%index] : vector<[4]x[4]xi1> from vector<4x[4]x[4]xi1>403 return %extract : vector<[4]x[4]xi1>404}405 406// -----407 408// CHECK-LABEL: @lift_illegal_transpose_to_memory(409// CHECK-SAME: %[[INDEXA:[a-z0-9]+]]: index,410// CHECK-SAME: %[[INDEXB:[a-z0-9]+]]: index,411// CHECK-SAME: %[[MEMREF:[a-z0-9]+]]: memref<?x?xf32>)412func.func @lift_illegal_transpose_to_memory(%a: index, %b: index, %memref: memref<?x?xf32>) -> vector<4x[8]xf32> {413 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index414 // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index415 // CHECK-DAG: %[[C0_F32:.*]] = arith.constant 0.000000e+00 : f32416 // CHECK-DAG: %[[VSCALE:.*]] = vector.vscale417 // CHECK-DAG: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index418 // CHECK-NEXT: %[[READ_SUBVIEW:.*]] = memref.subview %[[MEMREF]][%[[INDEXA]], %[[INDEXB]]] [%[[C8_VSCALE]], 4] [1, 1] : memref<?x?xf32> to memref<?x4xf32, strided<[?, 1], offset: ?>>419 // CHECK-NEXT: %[[CAST:.*]] = memref.cast %[[READ_SUBVIEW]] : memref<?x4xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>420 // CHECK-NEXT: %[[TRANSPOSE:.*]] = memref.transpose %[[CAST]] (d0, d1) -> (d1, d0) : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>421 // CHECK-NEXT: %[[LEGAL_READ:.*]] = vector.transfer_read %[[TRANSPOSE]][%c0, %c0], %[[C0_F32]] : memref<?x?xf32, strided<[?, ?], offset: ?>>, vector<4x[8]xf32>422 // CHECK-NEXT: return %[[LEGAL_READ]]423 %pad = arith.constant 0.0 : f32424 %illegalRead = vector.transfer_read %memref[%a, %b], %pad : memref<?x?xf32>, vector<[8]x4xf32>425 %legalType = vector.transpose %illegalRead, [1, 0] : vector<[8]x4xf32> to vector<4x[8]xf32>426 return %legalType : vector<4x[8]xf32>427}428 429// -----430 431// CHECK-LABEL: @lift_illegal_transpose_to_memory_with_mask(432// CHECK-SAME: %[[DIM0:[a-z0-9]+]]: index,433// CHECK-SAME: %[[DIM1:[a-z0-9]+]]: index,434// CHECK-SAME: %[[MEMREF:[a-z0-9]+]]: memref<?x?xf32>435func.func @lift_illegal_transpose_to_memory_with_mask(%dim0: index, %dim1: index, %memref: memref<?x?xf32>, %a: index, %b: index) -> vector<4x[8]xf32> {436 // CHECK-DAG: %[[READ_SUBVIEW:.*]] = memref.subview %[[MEMREF]]437 // CHECK-DAG: %[[CAST:.*]] = memref.cast %[[READ_SUBVIEW]]438 // CHECK-DAG: %[[TRANSPOSE:.*]] = memref.transpose %[[CAST]]439 // CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[DIM1]], %[[DIM0]] : vector<4x[8]xi1>440 // CHECK: %[[LEGAL_READ:.*]] = vector.transfer_read %[[TRANSPOSE]]441 // CHECK-SAME: %[[MASK]] : memref<?x?xf32, strided<[?, ?], offset: ?>>, vector<4x[8]xf32>442 // CHECK-NEXT: return %[[LEGAL_READ]]443 %pad = arith.constant 0.0 : f32444 %mask = vector.create_mask %dim0, %dim1 : vector<[8]x4xi1>445 %illegalRead = vector.transfer_read %memref[%a, %b], %pad, %mask : memref<?x?xf32>, vector<[8]x4xf32>446 %legalType = vector.transpose %illegalRead, [1, 0] : vector<[8]x4xf32> to vector<4x[8]xf32>447 return %legalType : vector<4x[8]xf32>448}449 450// -----451 452// CHECK-LABEL: @lift_illegal_transpose_to_memory_with_arith_extop(453// CHECK-SAME: %[[MEMREF:[a-z0-9]+]]: memref<?x?xi8>454func.func @lift_illegal_transpose_to_memory_with_arith_extop(%a: index, %b: index, %memref: memref<?x?xi8>) -> vector<4x[8]xi32> {455 // CHECK-DAG: %[[READ_SUBVIEW:.*]] = memref.subview %[[MEMREF]]456 // CHECK-DAG: %[[CAST:.*]] = memref.cast %[[READ_SUBVIEW]]457 // CHECK-DAG: %[[TRANSPOSE:.*]] = memref.transpose %[[CAST]]458 // CHECK: %[[LEGAL_READ:.*]] = vector.transfer_read %[[TRANSPOSE]]459 // CHECK-NEXT: %[[EXT_TYPE:.*]] = arith.extsi %[[LEGAL_READ]] : vector<4x[8]xi8> to vector<4x[8]xi32>460 // CHECK-NEXT: return %[[EXT_TYPE]]461 %pad = arith.constant 0 : i8462 %illegalRead = vector.transfer_read %memref[%a, %b], %pad : memref<?x?xi8>, vector<[8]x4xi8>463 %extRead = arith.extsi %illegalRead : vector<[8]x4xi8> to vector<[8]x4xi32>464 %legalType = vector.transpose %extRead, [1, 0] : vector<[8]x4xi32> to vector<4x[8]xi32>465 return %legalType : vector<4x[8]xi32>466}467 468// -----469 470// CHECK-LABEL: @lift_illegal_transpose_to_memory_with_in_bounds_attr471func.func @lift_illegal_transpose_to_memory_with_in_bounds_attr(%a: index, %b: index, %memref: memref<?x?xf32>) -> vector<4x[8]xf32> {472 // CHECK: vector.transfer_read473 // CHECK-SAME: in_bounds = [true, false]474 // CHECK-NOT: in_bounds = [false, true]475 %pad = arith.constant 0.0 : f32476 %illegalRead = vector.transfer_read %memref[%a, %b], %pad {in_bounds = [false, true]}: memref<?x?xf32>, vector<[8]x4xf32>477 %legalType = vector.transpose %illegalRead, [1, 0] : vector<[8]x4xf32> to vector<4x[8]xf32>478 return %legalType : vector<4x[8]xf32>479}480 481// -----482 483// The pass should do nothing (and not crash).484// CHECK-LABEL: @illegal_transpose_no_defining_source_op485func.func @illegal_transpose_no_defining_source_op(%vec: vector<[4]x1xf32>) -> vector<1x[4]xf32>486{487 // CHECK: vector.transpose488 %0 = vector.transpose %vec, [1, 0] : vector<[4]x1xf32> to vector<1x[4]xf32>489 return %0 : vector<1x[4]xf32>490}491 492// -----493 494// CHECK-LABEL: @multi_tile_splat495func.func @multi_tile_splat() -> vector<[8]x[8]xi32>496{497 // CHECK: %[[SPLAT:.*]] = arith.constant dense<42> : vector<[4]x[4]xi32>498 // CHECK-NEXT: return %[[SPLAT]], %[[SPLAT]], %[[SPLAT]], %[[SPLAT]] : vector<[4]x[4]xi32>, vector<[4]x[4]xi32>, vector<[4]x[4]xi32>, vector<[4]x[4]xi32>499 %0 = arith.constant dense<42> : vector<[8]x[8]xi32>500 return %0 : vector<[8]x[8]xi32>501}502 503// -----504 505// CHECK: #[[$TRANSPOSE_MAP_0:.*]] = affine_map<(d0, d1) -> (d1, d0)>506 507// CHECK-LABEL: @transpose_store_scalable_via_za(508// CHECK-SAME: %[[VEC:.*]]: vector<2x[4]xf32>509// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>,510// CHECK-SAME: %[[I:.*]]: index,511// CHECK-SAME: %[[J:.*]]: index)512func.func @transpose_store_scalable_via_za(%vec: vector<2x[4]xf32>, %dest: memref<?x?xf32>, %i: index, %j: index) {513 // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index514 // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index515 // CHECK-NEXT: %[[INIT:.*]] = arm_sme.get_tile : vector<[4]x[4]xf32>516 // CHECK-NEXT: %[[V0:.*]] = vector.extract %[[VEC]][0] : vector<[4]xf32> from vector<2x[4]xf32>517 // CHECK-NEXT: %[[R0:.*]] = vector.insert %[[V0]], %[[INIT]] [0] : vector<[4]xf32> into vector<[4]x[4]xf32>518 // CHECK-NEXT: %[[V1:.*]] = vector.extract %[[VEC]][1] : vector<[4]xf32> from vector<2x[4]xf32>519 // CHECK-NEXT: %[[RES:.*]] = vector.insert %[[V1]], %[[R0]] [1] : vector<[4]xf32> into vector<[4]x[4]xf32>520 // CHECK-NEXT: %[[VSCALE:.*]] = vector.vscale521 // CHECK-NEXT: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index522 // CHECK-NEXT: %[[MASK:.*]] = vector.create_mask %[[C4_VSCALE]], %[[C2]] : vector<[4]x[4]xi1>523 // CHECK-NEXT: vector.transfer_write %[[RES]], %[[DEST]][%[[I]], %[[J]]], %[[MASK]] {in_bounds = [true, true], permutation_map = #[[$TRANSPOSE_MAP_0]]} : vector<[4]x[4]xf32>, memref<?x?xf32>524 %tr = vector.transpose %vec, [1, 0] : vector<2x[4]xf32> to vector<[4]x2xf32>525 vector.transfer_write %tr, %dest[%i, %j] {in_bounds = [true, true]} : vector<[4]x2xf32>, memref<?x?xf32>526 return527}528 529// -----530 531// CHECK-LABEL: @transpose_store_scalable_via_za_masked(532// CHECK-SAME: %[[A:[a-z0-9]+]]: index,533// CHECK-SAME: %[[B:[a-z0-9]+]]: index)534func.func @transpose_store_scalable_via_za_masked(%vec: vector<2x[4]xf32>, %dest: memref<?x?xf32>, %a: index, %b: index) {535 // CHECK: %[[C2:.*]] = arith.constant 2 : index536 // CHECK: %[[MIN:.*]] = index.mins %[[B]], %[[C2]]537 // CHECK: %[[MASK:.*]] = vector.create_mask %[[A]], %[[MIN]] : vector<[4]x[4]xi1>538 // CHECK: vector.transfer_write {{.*}} %[[MASK]] {{.*}} : vector<[4]x[4]xf32>, memref<?x?xf32>539 %c0 = arith.constant 0 : index540 %mask = vector.create_mask %a, %b : vector<[4]x2xi1>541 %tr = vector.transpose %vec, [1, 0] : vector<2x[4]xf32> to vector<[4]x2xf32>542 vector.transfer_write %tr, %dest[%c0, %c0], %mask {in_bounds = [true, true]} : vector<[4]x2xf32>, memref<?x?xf32>543 return544}545 546// -----547 548// CHECK-LABEL: @transpose_store_scalable_via_za_multi_tile(549// CHECK-SAME: %[[VEC:.*]]: vector<8x[4]xf32>550// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>,551// CHECK-SAME: %[[I:.*]]: index,552// CHECK-SAME: %[[J:.*]]: index)553func.func @transpose_store_scalable_via_za_multi_tile(%vec: vector<8x[4]xf32>, %dest: memref<?x?xf32>, %i: index, %j: index) {554 // CHECK: %[[C4:.*]] = arith.constant 4 : index555 556 // <skip 3x other extract+insert chain>557 // CHECK: %[[V3:.*]] = vector.extract %[[VEC]][3] : vector<[4]xf32> from vector<8x[4]xf32>558 // CHECK: %[[TILE_0:.*]] = vector.insert %[[V3]], %{{.*}} [3] : vector<[4]xf32> into vector<[4]x[4]xf32>559 // CHECK: %[[VSCALE:.*]] = vector.vscale560 // CHECK: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index561 // CHECK: %[[MASK:.*]] = vector.create_mask %c4_vscale, %c4 : vector<[4]x[4]xi1>562 // CHECK: vector.transfer_write %[[TILE_0]], %[[DEST]][%[[I]], %[[J]]], %[[MASK]] {{.*}} : vector<[4]x[4]xf32>, memref<?x?xf32>563 564 // <skip 3x other extract+insert chain>565 // CHECK: %[[V7:.*]] = vector.extract %arg0[7] : vector<[4]xf32> from vector<8x[4]xf32>566 // CHECK: %[[TILE_1:.*]] = vector.insert %[[V7]], %{{.*}} [3] : vector<[4]xf32> into vector<[4]x[4]xf32>567 // CHECK: %[[J_OFFSET:.*]] = arith.addi %[[J]], %[[C4]] : index568 // CHECK: vector.transfer_write %[[TILE_1]], %[[DEST]][%[[I]], %[[J_OFFSET]]], %[[MASK]] {{.*}} : vector<[4]x[4]xf32>, memref<?x?xf32>569 %tr = vector.transpose %vec, [1, 0] : vector<8x[4]xf32> to vector<[4]x8xf32>570 vector.transfer_write %tr, %dest[%i, %j] {in_bounds = [true, true]} : vector<[4]x8xf32>, memref<?x?xf32>571 return572}573 574// -----575 576// CHECK-LABEL: @transpose_store_scalable_via_za_multi_tile_wide577func.func @transpose_store_scalable_via_za_multi_tile_wide(%vec: vector<2x[8]xf32>, %dest: memref<?x?xf32>, %i: index, %j: index) {578 // <check extracts from lower 4 x vscale of %vec>579 // CHECK: vector.scalable.extract580 // CHECK: %[[ROW_2_LOWER:.*]] = vector.scalable.extract %{{.*}}[0] : vector<[4]xf32> from vector<[8]xf32>581 // CHECK: %[[TILE_0:.*]] = vector.insert %[[ROW_2_LOWER]], %{{.*}}[1] : vector<[4]xf32> into vector<[4]x[4]xf32>582 // CHECK: vector.transfer_write %[[TILE_0]], %{{.*}}[%[[I:.[a-z0-9]+]], %[[J:[a-z0-9]+]]]583 584 // <check extracts from upper 4 x vscale of %vec>585 // CHECK: vector.scalable.extract586 // CHECK: %[[ROW_2_UPPER:.*]] = vector.scalable.extract %{{.*}}[4] : vector<[4]xf32> from vector<[8]xf32>587 // CHECK: %[[TILE_0:.*]] = vector.insert %[[ROW_2_UPPER]], %{{.*}}[1] : vector<[4]xf32> into vector<[4]x[4]xf32>588 // CHECK: %[[I_OFFSET:.*]] = arith.addi %c4_vscale, %[[I]] : index589 // CHECK: vector.transfer_write %[[TILE_0]], %{{.*}}[%[[I_OFFSET]], %[[J]]]590 %tr = vector.transpose %vec, [1, 0] : vector<2x[8]xf32> to vector<[8]x2xf32>591 vector.transfer_write %tr, %dest[%i, %j] {in_bounds = [true, true]} : vector<[8]x2xf32>, memref<?x?xf32>592 return593}594 595// -----596 597// CHECK-LABEL: @negative_transpose_store_scalable_via_za__bad_source_shape598// CHECK-NOT: arm_sme.get_tile599func.func @negative_transpose_store_scalable_via_za__bad_source_shape(%vec: vector<2x[7]xf32>, %dest: memref<?x?xf32>, %i: index, %j: index) {600 %tr = vector.transpose %vec, [1, 0] : vector<2x[7]xf32> to vector<[7]x2xf32>601 vector.transfer_write %tr, %dest[%i, %j] {in_bounds = [true, true]} : vector<[7]x2xf32>, memref<?x?xf32>602 return603}604 605// -----606 607// From: https://github.com/llvm/llvm-project/issues/118449.608// Check -arm-sme-vector-legalization does not crash when it encounters a `vector.mask` that609// does not contain a maskable op.610func.func @vector_mask_without_maskable_op(%mask: vector<16x2xi1>, %vec: vector<16x16xf32>) -> vector<16x16xf32> {611 %0 = vector.mask %mask { vector.yield %vec : vector<16x16xf32> } : vector<16x2xi1> -> vector<16x16xf32>612 return %0 : vector<16x16xf32>613}614 615// -----616 617//=============================================================================618// 1D examples - to be moved to the SVE dialect619//=============================================================================620 621/// TODO: Handle in_bounds622 623// CHECK-LABEL: func.func @xfer_read_scalable_column(624// CHECK-SAME: %[[IDX_0:[a-zA-Z0-9]+]]: index,625// CHECK-SAME: %[[IDX_1:[a-zA-Z0-9]+]]: index,626// CHECK-SAME: %[[PAD:.*]]: f32,627// CHECK-SAME: %[[SRC:.*]]: memref<?x?xf32>) -> vector<[4]x1xf32> {628func.func @xfer_read_scalable_column(%a: index, %b: index, %pad: f32, %src: memref<?x?xf32>) -> (vector<[4]x1xf32>) {629 // CHECK: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>630 // CHECK: %[[STEP:.*]] = arith.constant 1 : index631 // CHECK: %[[C4:.*]] = arith.constant 4 : index632 // CHECK: %[[LB:.*]] = arith.constant 0 : index633 // CHECK: %[[VSCALE:.*]] = vector.vscale634 // CHECK: %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index635 636 // <scf.for>637 // CHECK: %[[SCF:.*]] = scf.for %[[IND_VAR:.*]] = %[[LB]] to %[[C4_VSCALE]] step %[[STEP]] iter_args(%[[SCF_RES:.*]] = %[[INIT]]) -> (vector<[4]xf32>) {638 // CHECK: %[[IDX_0_UPDATED:.*]] = arith.addi %[[IND_VAR]], %[[IDX_0]] : index639 // CHECK: %[[VAL_10:.*]] = memref.load %[[SRC]][%[[IDX_0_UPDATED]], %[[IDX_1]]] : memref<?x?xf32>640 // CHECK: %[[RES_UPDATED:.*]] = vector.insert %[[VAL_10]], %[[SCF_RES]] [%[[IND_VAR]]] : f32 into vector<[4]xf32>641 // CHECK: scf.yield %[[RES_UPDATED]] : vector<[4]xf32>642 // CHECK: }643 644 // <shape-cast>645 // CHECK: %[[SC:.*]] = vector.shape_cast %[[SCF]] : vector<[4]xf32> to vector<[4]x1xf32>646 // CHECK: return %[[SC]]647 %read = vector.transfer_read %src[%a, %b], %pad : memref<?x?xf32>, vector<[4]x1xf32>648 return %read : vector<[4]x1xf32>649}650 651// -----652 653// CHECK-LABEL: func.func @negative_xfer_read_scalable_column_x2654func.func @negative_xfer_read_scalable_column_x2(%a: index, %b: index, %pad: f32, %src: memref<?x?xf32>) -> (vector<[4]x2xf32>) {655 // CHECK-NOT: scf.for656 // CHECK-NOT: memref.load657 %read = vector.transfer_read %src[%a, %b], %pad : memref<?x?xf32>, vector<[4]x2xf32>658 return %read : vector<[4]x2xf32>659}660 661// -----662 663// CHECK-LABEL: func.func @negative_xfer_read_scalable_column_scalable_trailing_dim664func.func @negative_xfer_read_scalable_column_scalable_trailing_dim(%a: index, %b: index, %pad: f32, %src: memref<?x?xf32>) -> (vector<4x[1]xf32>) {665 // CHECK-NOT: scf.for666 // CHECK-NOT: memref.load667 %read = vector.transfer_read %src[%a, %b], %pad : memref<?x?xf32>, vector<4x[1]xf32>668 return %read : vector<4x[1]xf32>669}670