126 lines · plain
1// RUN: mlir-opt %s --test-vector-emulate-masked-load-store | FileCheck %s2 3// CHECK-LABEL: @vector_maskedload4// CHECK-SAME: (%[[ARG0:.*]]: memref<4x5xf32>) -> vector<4xf32> {5// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>6// CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index7// CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index8// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index9// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index10// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index11// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index12// CHECK-DAG: %[[S0:.*]] = vector.create_mask %[[C1]] : vector<4xi1>13// CHECK: %[[S1:.*]] = vector.extract %[[S0]][0] : i1 from vector<4xi1>14// CHECK: %[[S2:.*]] = scf.if %[[S1]] -> (vector<4xf32>) {15// CHECK: %[[S9:.*]] = memref.load %[[ARG0]][%[[C0]], %[[C4]]] : memref<4x5xf32>16// CHECK: %[[S10:.*]] = vector.insert %[[S9]], %[[CST]] [0] : f32 into vector<4xf32>17// CHECK: scf.yield %[[S10]] : vector<4xf32>18// CHECK: } else {19// CHECK: scf.yield %[[CST]] : vector<4xf32>20// CHECK: }21// CHECK: %[[S3:.*]] = vector.extract %[[S0]][1] : i1 from vector<4xi1>22// CHECK: %[[S4:.*]] = scf.if %[[S3]] -> (vector<4xf32>) {23// CHECK: %[[S9:.*]] = memref.load %[[ARG0]][%[[C0]], %[[C5]]] : memref<4x5xf32>24// CHECK: %[[S10:.*]] = vector.insert %[[S9]], %[[S2]] [1] : f32 into vector<4xf32>25// CHECK: scf.yield %[[S10]] : vector<4xf32>26// CHECK: } else {27// CHECK: scf.yield %[[S2]] : vector<4xf32>28// CHECK: }29// CHECK: %[[S5:.*]] = vector.extract %[[S0]][2] : i1 from vector<4xi1>30// CHECK: %[[S6:.*]] = scf.if %[[S5]] -> (vector<4xf32>) {31// CHECK: %[[S9:.*]] = memref.load %[[ARG0]][%[[C0]], %[[C6]]] : memref<4x5xf32>32// CHECK: %[[S10:.*]] = vector.insert %[[S9]], %[[S4]] [2] : f32 into vector<4xf32>33// CHECK: scf.yield %[[S10]] : vector<4xf32>34// CHECK: } else {35// CHECK: scf.yield %[[S4]] : vector<4xf32>36// CHECK: }37// CHECK: %[[S7:.*]] = vector.extract %[[S0]][3] : i1 from vector<4xi1>38// CHECK: %[[S8:.*]] = scf.if %[[S7]] -> (vector<4xf32>) {39// CHECK: %[[S9:.*]] = memref.load %[[ARG0]][%[[C0]], %[[C7]]] : memref<4x5xf32>40// CHECK: %[[S10:.*]] = vector.insert %[[S9]], %[[S6]] [3] : f32 into vector<4xf32>41// CHECK: scf.yield %[[S10]] : vector<4xf32>42// CHECK: } else {43// CHECK: scf.yield %[[S6]] : vector<4xf32>44// CHECK: }45// CHECK: return %[[S8]] : vector<4xf32>46func.func @vector_maskedload(%arg0 : memref<4x5xf32>) -> vector<4xf32> {47 %idx_0 = arith.constant 0 : index48 %idx_1 = arith.constant 1 : index49 %idx_4 = arith.constant 4 : index50 %mask = vector.create_mask %idx_1 : vector<4xi1>51 %s = arith.constant 0.0 : f3252 %pass_thru = vector.broadcast %s : f32 to vector<4xf32>53 %0 = vector.maskedload %arg0[%idx_0, %idx_4], %mask, %pass_thru : memref<4x5xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>54 return %0: vector<4xf32>55}56 57// CHECK-LABEL: @vector_maskedload_with_alignment58// CHECK: memref.load59// CHECK-SAME: {alignment = 8 : i64}60// CHECK: memref.load61// CHECK-SAME: {alignment = 8 : i64}62func.func @vector_maskedload_with_alignment(%arg0 : memref<4x5xf32>) -> vector<4xf32> {63 %idx_0 = arith.constant 0 : index64 %idx_1 = arith.constant 1 : index65 %idx_4 = arith.constant 4 : index66 %mask = vector.create_mask %idx_1 : vector<4xi1>67 %s = arith.constant 0.0 : f3268 %pass_thru = vector.broadcast %s : f32 to vector<4xf32>69 %0 = vector.maskedload %arg0[%idx_0, %idx_4], %mask, %pass_thru {alignment = 8}: memref<4x5xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>70 return %0: vector<4xf32>71}72 73// CHECK-LABEL: @vector_maskedstore74// CHECK-SAME: (%[[ARG0:.*]]: memref<4x5xf32>, %[[ARG1:.*]]: vector<4xf32>) {75// CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index76// CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index77// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index78// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index79// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index80// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index81// CHECK-DAG: %[[S0:.*]] = vector.create_mask %[[C1]] : vector<4xi1>82// CHECK: %[[S1:.*]] = vector.extract %[[S0]][0] : i1 from vector<4xi1>83// CHECK: scf.if %[[S1]] {84// CHECK: %[[S5:.*]] = vector.extract %[[ARG1]][0] : f32 from vector<4xf32>85// CHECK: memref.store %[[S5]], %[[ARG0]][%[[C0]], %[[C4]]] : memref<4x5xf32>86// CHECK: }87// CHECK: %[[S2:.*]] = vector.extract %[[S0]][1] : i1 from vector<4xi1>88// CHECK: scf.if %[[S2]] {89// CHECK: %[[S5:.*]] = vector.extract %[[ARG1]][1] : f32 from vector<4xf32>90// CHECK: memref.store %[[S5]], %[[ARG0]][%[[C0]], %[[C5]]] : memref<4x5xf32>91// CHECK: }92// CHECK: %[[S3:.*]] = vector.extract %[[S0]][2] : i1 from vector<4xi1>93// CHECK: scf.if %[[S3]] {94// CHECK: %[[S5:.*]] = vector.extract %[[ARG1]][2] : f32 from vector<4xf32>95// CHECK: memref.store %[[S5]], %[[ARG0]][%[[C0]], %[[C6]]] : memref<4x5xf32>96// CHECK: }97// CHECK: %[[S4:.*]] = vector.extract %[[S0]][3] : i1 from vector<4xi1>98// CHECK: scf.if %[[S4]] {99// CHECK: %[[S5:.*]] = vector.extract %[[ARG1]][3] : f32 from vector<4xf32>100// CHECK: memref.store %[[S5]], %[[ARG0]][%[[C0]], %[[C7]]] : memref<4x5xf32>101// CHECK: }102// CHECK: return103// CHECK:}104func.func @vector_maskedstore(%arg0 : memref<4x5xf32>, %arg1 : vector<4xf32>) {105 %idx_0 = arith.constant 0 : index106 %idx_1 = arith.constant 1 : index107 %idx_4 = arith.constant 4 : index108 %mask = vector.create_mask %idx_1 : vector<4xi1>109 vector.maskedstore %arg0[%idx_0, %idx_4], %mask, %arg1 : memref<4x5xf32>, vector<4xi1>, vector<4xf32>110 return111}112 113// CHECK-LABEL: @vector_maskedstore_with_alignment114// CHECK: memref.store115// CHECK-SAME: {alignment = 8 : i64}116// CHECK: memref.store117// CHECK-SAME: {alignment = 8 : i64}118func.func @vector_maskedstore_with_alignment(%arg0 : memref<4x5xf32>, %arg1 : vector<4xf32>) {119 %idx_0 = arith.constant 0 : index120 %idx_1 = arith.constant 1 : index121 %idx_4 = arith.constant 4 : index122 %mask = vector.create_mask %idx_1 : vector<4xi1>123 vector.maskedstore %arg0[%idx_0, %idx_4], %mask, %arg1 { alignment = 8 } : memref<4x5xf32>, vector<4xi1>, vector<4xf32>124 return125}126