47 lines · plain
1// RUN: mlir-opt %s -transform-preload-library='transform-library-paths=%p/td/unroll-elements.mlir' \2// RUN: -transform-interpreter=entry-point=unroll_to_elements | FileCheck %s3 4//===----------------------------------------------------------------------===//5// Test UnrollFromElements.6//===----------------------------------------------------------------------===//7 8// CHECK-LABEL: @unroll_from_elements_2d9// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32)10// CHECK-NEXT: %[[UNDEF_RES:.*]] = ub.poison : vector<2x2xf32>11// CHECK-NEXT: %[[VEC_0:.*]] = vector.from_elements %[[ARG0]], %[[ARG1]] : vector<2xf32>12// CHECK-NEXT: %[[RES_0:.*]] = vector.insert %[[VEC_0]], %[[UNDEF_RES]] [0] : vector<2xf32> into vector<2x2xf32>13// CHECK-NEXT: %[[VEC_1:.*]] = vector.from_elements %[[ARG2]], %[[ARG3]] : vector<2xf32>14// CHECK-NEXT: %[[RES_1:.*]] = vector.insert %[[VEC_1]], %[[RES_0]] [1] : vector<2xf32> into vector<2x2xf32>15// CHECK-NEXT: return %[[RES_1]] : vector<2x2xf32>16func.func @unroll_from_elements_2d(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32) -> vector<2x2xf32> {17 %0 = vector.from_elements %arg0, %arg1, %arg2, %arg3 : vector<2x2xf32>18 return %0 : vector<2x2xf32>19}20 21// CHECK-LABEL: @unroll_from_elements_3d22// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32)23// CHECK-NEXT: %[[UNDEF_RES:.*]] = ub.poison : vector<2x1x2xf32>24// CHECK-NEXT: %[[UNDEF_RANK_2:.*]] = ub.poison : vector<1x2xf32>25// CHECK-NEXT: %[[VEC_0:.*]] = vector.from_elements %[[ARG0]], %[[ARG1]] : vector<2xf32>26// CHECK-NEXT: %[[RANK_2_0:.*]] = vector.insert %[[VEC_0]], %[[UNDEF_RANK_2]] [0] : vector<2xf32> into vector<1x2xf32>27// CHECK-NEXT: %[[RES_0:.*]] = vector.insert %[[RANK_2_0]], %[[UNDEF_RES]] [0] : vector<1x2xf32> into vector<2x1x2xf32>28// CHECK-NEXT: %[[VEC_1:.*]] = vector.from_elements %[[ARG2]], %[[ARG3]] : vector<2xf32>29// CHECK-NEXT: %[[RANK_2_1:.*]] = vector.insert %[[VEC_1]], %[[UNDEF_RANK_2]] [0] : vector<2xf32> into vector<1x2xf32>30// CHECK-NEXT: %[[RES_1:.*]] = vector.insert %[[RANK_2_1]], %[[RES_0]] [1] : vector<1x2xf32> into vector<2x1x2xf32>31// CHECK-NEXT: return %[[RES_1]] : vector<2x1x2xf32>32func.func @unroll_from_elements_3d(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32) -> vector<2x1x2xf32> {33 %0 = vector.from_elements %arg0, %arg1, %arg2, %arg3 : vector<2x1x2xf32>34 return %0 : vector<2x1x2xf32>35}36 37// 1-D vector.from_elements should not be unrolled.38 39// CHECK-LABEL: @negative_unroll_from_elements_1d40// CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32)41// CHECK-NEXT: %[[RES:.*]] = vector.from_elements %[[ARG0]], %[[ARG1]] : vector<2xf32>42// CHECK-NEXT: return %[[RES]] : vector<2xf32>43func.func @negative_unroll_from_elements_1d(%arg0: f32, %arg1: f32) -> vector<2xf32> {44 %0 = vector.from_elements %arg0, %arg1 : vector<2xf32>45 return %0 : vector<2xf32>46}47