115 lines · plain
1// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-expand-shape-bubbling %s | FileCheck %s2 3func.func @bubble_parallel_reshapes(%arg0: tensor<?x?x?x?xf32>, %s0: index, %s1: index, %s2: index, %s3: index) -> tensor<?x?x?x?xf32> {4 %collapse = tensor.collapse_shape %arg0 [[0], [1, 2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>5 %expand = tensor.expand_shape %collapse [[0], [1], [2, 3]]6 output_shape [%s0, %s1, %s2, %s3] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>7 return %expand : tensor<?x?x?x?xf32>8}9// CHECK: func @bubble_parallel_reshapes10// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>11// CHECK-SAME: %[[S0:.+]]: index, %[[S1:.+]]: index, %[[S2:.+]]: index, %[[S3:.+]]: index12// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index13// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index14// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?x?x?xf32>15// CHECK-DAG: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<?x?x?x?xf32>16// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1], [2], [3, 4]]17// CHECK-SAME: output_shape [%[[S0]], %[[DIM1]], %[[DIM2]], %[[S2]], %[[S3]]] : tensor<?x?x?x?xf32> into tensor<?x?x?x?x?xf32>18// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[EXPAND]] {{\[}}[0], [1, 2], [3], [4]] : tensor<?x?x?x?x?xf32> into tensor<?x?x?x?xf32>19// CHECK: return %[[COLLAPSE]]20 21// -----22 23func.func @no_bubble_full_intersecting_reshapes(%arg0: tensor<?x?x?x?xf32>, %s0: index, %s1: index, %s2: index, %s3: index) -> tensor<?x?x?x?xf32> {24 %collapse = tensor.collapse_shape %arg0 [[0], [1, 2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>25 %expand = tensor.expand_shape %collapse [[0], [1, 2], [3]]26 output_shape [%s0, %s1, %s2, %s3] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>27 return %expand : tensor<?x?x?x?xf32>28}29// CHECK: func @no_bubble_full_intersecting_reshapes30// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>31// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3]]32// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0], [1, 2], [3]]33// CHECK: return %[[EXPAND]]34 35// -----36 37func.func @no_bubble_partial_intersecting_reshapes(%arg0: tensor<?x?x?x?xf32>, %s0: index, %s1: index, %s2: index, %s3: index) -> tensor<?x?x?x?xf32> {38 %collapse = tensor.collapse_shape %arg0 [[0, 1, 2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?xf32>39 %expand = tensor.expand_shape %collapse [[0, 1], [2, 3]]40 output_shape [%s0, %s1, %s2, %s3] : tensor<?x?xf32> into tensor<?x?x?x?xf32>41 return %expand : tensor<?x?x?x?xf32>42}43// CHECK: func @no_bubble_partial_intersecting_reshapes44// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>45// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]46// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0, 1], [2, 3]]47// CHECK: return %[[EXPAND]]48 49// -----50 51func.func @no_bubble_0d_tensor_reshapes(%arg0: tensor<1x1xf32>) -> tensor<1x1x1xf32> {52 %collapse = tensor.collapse_shape %arg0 [] : tensor<1x1xf32> into tensor<f32>53 %expand = tensor.expand_shape %collapse []54 output_shape [1, 1, 1] : tensor<f32> into tensor<1x1x1xf32>55 return %expand : tensor<1x1x1xf32>56}57// CHECK: func @no_bubble_0d_tensor_reshapes58// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1xf32>59// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}]60// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}]61// CHECK: return %[[EXPAND]]62 63// -----64 65// Test the case where the reassocation indices in the collapse and expand66// are of same size.67func.func @bubble_expand_match_non_unit_size_reassocation(68 %arg0 : tensor<4x?x4x32x4x?xf16>, %arg1 : index, %arg2 : index) -> tensor<4x?x4x128x?x32xf16> {69 %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2], [3, 4], [5]]70 : tensor<4x?x4x32x4x?xf16> into tensor<?x128x?xf16>71 %expanded = tensor.expand_shape %collapsed [[0, 1, 2], [3], [4, 5]] output_shape [4, %arg1, 4, 128, %arg2, 32]72 : tensor<?x128x?xf16> into tensor<4x?x4x128x?x32xf16>73 return %expanded : tensor<4x?x4x128x?x32xf16>74}75// CHECK: func @bubble_expand_match_non_unit_size_reassocation76// CHECK-SAME: %[[ARG0:.+]]: tensor<4x?x4x32x4x?xf16>77// CHECK-SAME: %[[ARG1:[a-zA-z0-9]+]]: index78// CHECK-SAME: %[[ARG2:[a-zA-z0-9]+]]: index79// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]]80// CHECK-SAME: {{\[}}[0], [1], [2], [3], [4], [5, 6]{{\]}}81// CHECK-SAME: [4, %[[ARG1]], 4, 32, 4, %[[ARG2]], 32]82// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[EXPANDED]]83// CHECK-SAME: {{\[}}[0], [1], [2], [3, 4], [5], [6]{{\]}}84// CHECK: return %[[COLLAPSED]]85 86// -----87 88// Test the case where the trailing collapse isnt needed.89func.func @no_collapse_generated(90 %arg0 : tensor<4x?x4x128x?xf16>, %arg1 : index, %arg2 : index) -> tensor<4x?x4x128x?x32xf16> {91 %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2], [3], [4]]92 : tensor<4x?x4x128x?xf16> into tensor<?x128x?xf16>93 %expanded = tensor.expand_shape %collapsed [[0, 1, 2], [3], [4, 5]] output_shape [4, %arg1, 4, 128, %arg2, 32]94 : tensor<?x128x?xf16> into tensor<4x?x4x128x?x32xf16>95 return %expanded : tensor<4x?x4x128x?x32xf16>96}97// CHECK: func @no_collapse_generated98// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape 99// CHECK: return %[[EXPANDED]]100 101// -----102 103// Test the case where the leading expand isnt needed.104func.func @no_expand_generated(105 %arg0 : tensor<4x?x4x128x?x?x?xf16>, %arg1 : index, %arg2 : index, %arg3 : index) -> tensor<4x?x4x128x?x?xf16> {106 %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2], [3], [4], [5, 6]]107 : tensor<4x?x4x128x?x?x?xf16> into tensor<?x128x?x?xf16>108 %expanded = tensor.expand_shape %collapsed [[0, 1, 2], [3], [4], [5]] output_shape [4, %arg1, 4, 128, %arg2, %arg3]109 : tensor<?x128x?x?xf16> into tensor<4x?x4x128x?x?xf16>110 return %expanded : tensor<4x?x4x128x?x?xf16>111}112// CHECK: func @no_expand_generated113// CHECK: %[[EXPANDED:.+]] = tensor.collapse_shape114// CHECK: return %[[EXPANDED]]115