130 lines · plain
1// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing -split-input-file | FileCheck %s2 3// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>4// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1)>5 6// CHECK-LABEL: func @reshape7// CHECK-SAME: (%[[A:.*]]: tensor<?x16xf32>, %[[B:.*]]: tensor<16xf32>, %[[INIT:.*]]: tensor<?x112x16xf32>, %[[SZ0:.*]]: index)8// CHECK: %[[C0:.*]] = arith.constant 0 : index9// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[A]]10// CHECK: %[[DIM:.*]] = tensor.dim %[[EXPANDED]], %[[C0]]11// CHECK: %[[RI:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2]] : tensor<?x112x16xf32> into tensor<?x16xf32>12// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP3]], #[[$MAP2]]],13// CHECK-SAME: iterator_types = ["parallel", "parallel"]}14// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<?x16xf32>)15// CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[\[}}0, 1], [2]] output_shape [%[[DIM]], 112, 16] : tensor<?x16xf32> into tensor<?x112x16xf32>16// CHECK: return %[[RR]] : tensor<?x112x16xf32>17func.func @reshape(%A: tensor<?x16xf32>, %B: tensor<16xf32>, %init: tensor<?x112x16xf32>, %sz0: index) -> tensor<?x112x16xf32> {18 %0 = tensor.expand_shape %A [[0, 1], [2]] output_shape [%sz0, 112, 16]19 : tensor<?x16xf32> into tensor<?x112x16xf32>20 %2 = linalg.generic {indexing_maps = [21 affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d2)>,22 affine_map<(d0, d1, d2) -> (d0, d1, d2)>],23 iterator_types = ["parallel", "parallel", "parallel"]}24 ins(%0, %B : tensor<?x112x16xf32>, tensor<16xf32>)25 outs(%init : tensor<?x112x16xf32>) {26 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):27 %s = arith.subf %arg1, %arg2 : f3228 linalg.yield %s : f3229 } -> tensor<?x112x16xf32>30 return %2 : tensor<?x112x16xf32>31}32 33// -----34 35// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>36// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1)>37 38// CHECK-LABEL: func @reshape_multiple39// CHECK-SAME: (%[[A:.*]]: tensor<12544x16xf32>, %[[B:.*]]: tensor<12544x16xf32>, %[[C:.*]]: tensor<16xf32>)40// CHECK: %[[I:.*]] = tensor.empty() : tensor<112x112x16xf32>41// CHECK: %[[RI:.*]] = tensor.collapse_shape %[[I]] {{\[}}[0, 1], [2]] : tensor<112x112x16xf32> into tensor<12544x16xf32>42// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP2]], #[[$MAP3]], #[[$MAP2]]],43// CHECK-SAME: iterator_types = ["parallel", "parallel"]}44// CHECK-SAME: ins(%[[A]], %[[B]], %[[C]] : tensor<12544x16xf32>, tensor<12544x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<12544x16xf32>)45// CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] output_shape [112, 112, 16] : tensor<12544x16xf32> into tensor<112x112x16xf32>46// CHECK: return %[[RR]] : tensor<112x112x16xf32>47func.func @reshape_multiple(%A: tensor<12544x16xf32>, %B: tensor<12544x16xf32>,48 %C: tensor<16xf32>) -> tensor<112x112x16xf32> {49 %0 = tensor.expand_shape %A [[0, 1], [2]] output_shape [112, 112, 16]50 : tensor<12544x16xf32> into tensor<112x112x16xf32>51 %1 = tensor.expand_shape %B [[0, 1], [2]] output_shape [112, 112, 16]52 : tensor<12544x16xf32> into tensor<112x112x16xf32>53 %2 = tensor.empty() : tensor<112x112x16xf32>54 %3 = linalg.generic {indexing_maps = [55 affine_map<(d0, d1, d2) -> (d0, d1, d2)>,56 affine_map<(d0, d1, d2) -> (d0, d1, d2)>,57 affine_map<(d0, d1, d2) -> (d2)>,58 affine_map<(d0, d1, d2) -> (d0, d1, d2)>],59 iterator_types = ["parallel", "parallel", "parallel"]}60 ins(%0, %1, %C : tensor<112x112x16xf32>, tensor<112x112x16xf32>, tensor<16xf32>)61 outs(%2 : tensor<112x112x16xf32>) {62 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32):63 %s = arith.subf %arg1, %arg2 : f3264 %m = arith.mulf %s, %arg3 : f3265 linalg.yield %m : f3266 } -> tensor<112x112x16xf32>67 return %3 : tensor<112x112x16xf32>68}69 70// -----71 72// Negative test, since the second source is broadcasted from d1 we cannot merge73// d0 and d1 dimensions74// CHECK-LABEL: func @reshape_negative75// CHECK: tensor.expand_shape {{.*}} {{\[\[}}0, 1], [2]] output_shape [112, 112, 16] : tensor<12544x16xf32> into tensor<112x112x16xf32>76// CHECK: linalg.generic77// CHECK: } -> tensor<112x112x16xf32>78func.func @reshape_negative(%A: tensor<12544x16xf32>, %B: tensor<112xf32>) -> tensor<112x112x16xf32> {79 %20 = tensor.expand_shape %A [[0, 1], [2]] output_shape [112, 112, 16]80 : tensor<12544x16xf32> into tensor<112x112x16xf32>81 %21 = tensor.empty() : tensor<112x112x16xf32>82 %22 = linalg.generic {indexing_maps = [83 affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d1)>,84 affine_map<(d0, d1, d2) -> (d0, d1, d2)>],85 iterator_types = ["parallel", "parallel", "parallel"]}86 ins(%20, %B : tensor<112x112x16xf32>, tensor<112xf32>)87 outs(%21 : tensor<112x112x16xf32>) {88 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):89 %s = arith.subf %arg1, %arg2 : f3290 linalg.yield %s : f3291 } -> tensor<112x112x16xf32>92 return %22 : tensor<112x112x16xf32>93}94 95// -----96 97func.func @type_correctness(%arg0 : tensor<6x5xi32>, %arg1 : tensor<5xf32>,98 %arg2 : tensor<5xf32>) -> tensor<2x3x5xf32> {99 %cst_6 = arith.constant 1.000000e+00 : f32100 %cst_7 = arith.constant 7.000000e+00 : f32101 %cst_8 = arith.constant 1.1920929E-7 : f32102 %25 = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [2, 3, 5]103 : tensor<6x5xi32> into tensor<2x3x5xi32>104 %26 = tensor.empty() : tensor<2x3x5xf32>105 %28 = linalg.generic {106 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,107 affine_map<(d0, d1, d2) -> (d2)>,108 affine_map<(d0, d1, d2) -> (d2)>,109 affine_map<(d0, d1, d2) -> (d0, d1, d2)>],110 iterator_types = ["parallel", "parallel", "parallel"]}111 ins(%25, %arg1, %arg2 : tensor<2x3x5xi32>, tensor<5xf32>, tensor<5xf32>)112 outs(%26 : tensor<2x3x5xf32>) {113 ^bb0(%arg6: i32, %arg7: f32, %arg8: f32, %arg9: f32):114 %29 = arith.sitofp %arg6 : i32 to f32115 %30 = arith.addf %arg7, %cst_8 : f32116 %31 = arith.divf %cst_7, %30 : f32117 %32 = arith.divf %cst_6, %31 : f32118 %33 = arith.mulf %29, %32 : f32119 %34 = arith.addf %33, %arg8 : f32120 linalg.yield %34 : f32121 } -> tensor<2x3x5xf32>122 return %28 : tensor<2x3x5xf32>123}124// CHECK-LABEL: func @type_correctness125// CHECK: %[[OP:.+]] = linalg.generic126// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<6x5xi32>, tensor<5xf32>, tensor<5xf32>)127// CHECK-SAME: outs(%{{.+}} : tensor<6x5xf32>)128// CHECK: tensor.expand_shape %[[OP]]129// CHECK-SAME: tensor<6x5xf32> into tensor<2x3x5xf32>130