32 lines · plain
1// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-extract-from-collapse-shape %s | FileCheck %s2 3// CHECK-LABEL: @extract_from_collapse_shape4// CHECK-SAME: (%[[ARG0:.*]]: tensor<1x1x8xi8>)5func.func @extract_from_collapse_shape(%arg0: tensor<1x1x8xi8>) -> (i8, i8) {6 %c1 = arith.constant 1 : index7 %c0 = arith.constant 0 : index8 %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2]] : tensor<1x1x8xi8> into tensor<8xi8>9 %extracted = tensor.extract %collapsed[%c0] : tensor<8xi8>10 %extracted_0 = tensor.extract %collapsed[%c1] : tensor<8xi8>11 func.return %extracted, %extracted_0 : i8, i812}13 14// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index15// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index16// CHECK-DAG: %[[RESULT0:.*]] = tensor.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] : tensor<1x1x8xi8>17// CHECK-DAG: %[[RESULT1:.*]] = tensor.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C1]]] : tensor<1x1x8xi8>18// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]] : i8, i819 20// -----21 22// CHECK-LABEL: @extract_from_static_shape23// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x6x32xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)24func.func @extract_from_static_shape(%arg0 : tensor<2x6x32xf32>, %arg1 : index, %arg2 : index) -> f32 {25 %0 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<2x6x32xf32> into tensor<12x32xf32>26 %1 = tensor.extract %0[%arg1, %arg2] : tensor<12x32xf32>27 return %1 : f3228}29// CHECK-NEXT: %[[MODIFIED_INDEXES:.*]]:2 = affine.delinearize_index %[[ARG1]] into (2, 6)30// CHECK-NEXT: %[[RESULT:.*]] = tensor.extract %[[ARG0]][%[[MODIFIED_INDEXES]]#0, %[[MODIFIED_INDEXES]]#1, %[[ARG2]]] : tensor<2x6x32xf32>31// CHECK-NEXT: return %[[RESULT]] : f3232