218 lines · plain
1//RUN: mlir-opt -test-linalg-rank-reduce-contraction-ops --canonicalize -split-input-file %s | FileCheck %s2 3func.func @singleton_batch_matmul_tensor(%arg0 : tensor<1x128x512xf32>, %arg1 : tensor<1x512x256xf32>, %arg2: tensor<1x128x256xf32>) -> tensor<1x128x256xf32> {4 // CHECK-LABEL: @singleton_batch_matmul_tensor5 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<1x128x512xf32>6 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x512x256xf32>7 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<1x128x256xf32>8 // CHECK-NEXT: %[[COLLAPSED_LHS:.*]] = tensor.collapse_shape %[[LHS]] {{\[}}[0, 1], [2]]9 // CHECK-NEXT: %[[COLLAPSED_RHS:.*]] = tensor.collapse_shape %[[RHS]] {{\[}}[0, 1], [2]]10 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2]]11 // CHECK-NEXT: %[[MATMUL:.+]] = linalg.matmul ins(%[[COLLAPSED_LHS]], %[[COLLAPSED_RHS]] : tensor<128x512xf32>, tensor<512x256xf32>) outs(%[[COLLAPSED_INIT]] : tensor<128x256xf32>)12 // CHECK-NEXT: %[[RES:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2]] output_shape [1, 128, 256]13 // CHECK-NEXT: return %[[RES]]14 %1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x128x512xf32>, tensor<1x512x256xf32>)15 outs(%arg2 : tensor<1x128x256xf32>) -> tensor<1x128x256xf32>16 return %1 : tensor<1x128x256xf32>17}18 19// -----20 21func.func @singleton_batch_matmul_memref(%arg0 : memref<1x?x?xf32>, %arg1 : memref<1x?x?xf32>, %arg2: memref<1x?x?xf32>) {22 // CHECK-LABEL: @singleton_batch_matmul_memref23 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: memref<1x?x?xf32>24 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: memref<1x?x?xf32>25 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: memref<1x?x?xf32>26 // CHECK-NEXT: %[[COLLAPSED_LHS:.*]] = memref.collapse_shape %[[LHS]] {{\[}}[0, 1], [2]]27 // CHECK-NEXT: %[[COLLAPSED_RHS:.*]] = memref.collapse_shape %[[RHS]] {{\[}}[0, 1], [2]]28 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = memref.collapse_shape %[[INIT]] {{\[}}[0, 1], [2]]29 // CHECK-NEXT: linalg.matmul ins(%[[COLLAPSED_LHS]], %[[COLLAPSED_RHS]] : memref<?x?xf32>, memref<?x?xf32>) outs(%[[COLLAPSED_INIT]] : memref<?x?xf32>)30 // CHECK-NEXT: return31 linalg.batch_matmul ins(%arg0, %arg1 : memref<1x?x?xf32>, memref<1x?x?xf32>)32 outs(%arg2 : memref<1x?x?xf32>)33 return34}35 36// -----37 38func.func @negative_singleton_batch_matmul_to_matmul_memref(%arg0 : memref<1x?x?xf32>, %arg1 : memref<1x?x?xf32>, %arg2: memref<1x?x?xf32>) {39 // CHECK-LABEL: @negative_singleton_batch_matmul_to_matmul_memref40 // CHECK-NOT: collapse_shape41 // CHECK-NOT: linalg.matmul42 // CHECK-NOT: expand_shape43 linalg.batch_matmul indexing_maps = [44 affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,45 affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>,46 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>47 ]48 ins(%arg0, %arg1 : memref<1x?x?xf32>, memref<1x?x?xf32>)49 outs(%arg2 : memref<1x?x?xf32>)50 return51}52 53// -----54 55func.func @singleton_batch_matvec(%arg0 : tensor<1x128x512xf32>, %arg1 : tensor<1x512xf32>, %arg2: tensor<1x128xf32>) -> tensor<1x128xf32> {56 // CHECK-LABEL: @singleton_batch_matvec57 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<1x128x512xf32>58 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x512xf32>59 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<1x128xf32>60 // CHECK-NEXT: %[[COLLAPSED_LHS:.*]] = tensor.collapse_shape %[[LHS]] {{\[}}[0, 1], [2]]61 // CHECK-NEXT: %[[COLLAPSED_RHS:.*]] = tensor.collapse_shape %[[RHS]] {{\[}}[0, 1]]62 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1]]63 // CHECK-NEXT: %[[MATMUL:.+]] = linalg.matvec64 // CHECK-SAME: ins(%[[COLLAPSED_LHS]], %[[COLLAPSED_RHS]] : tensor<128x512xf32>, tensor<512xf32>) outs(%[[COLLAPSED_INIT]] : tensor<128xf32>)65 // CHECK-NEXT: %[[RES:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1]] output_shape [1, 128]66 // CHECK-NEXT: return %[[RES]]67 %1 = linalg.batch_matvec ins(%arg0, %arg1 : tensor<1x128x512xf32>, tensor<1x512xf32>)68 outs(%arg2 : tensor<1x128xf32>) -> tensor<1x128xf32>69 return %1 : tensor<1x128xf32>70}71 72// -----73 74func.func @singleton_batch_vecmat(%arg0 : tensor<1x?xf32>, %arg1 : tensor<1x?x?xf32>, %arg2: tensor<1x?xf32>) -> tensor<1x?xf32> {75 // CHECK-LABEL: @singleton_batch_vecmat76 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<1x?xf32>77 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x?x?xf32>78 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<1x?xf32>79 // CHECK-DAG: %[[C0:.*]] = arith.constant 080 // CHECK-NEXT: %[[COLLAPSED_LHS:.*]] = tensor.collapse_shape %[[LHS]] {{\[}}[0, 1]]81 // CHECK-NEXT: %[[COLLAPSED_RHS:.*]] = tensor.collapse_shape %[[RHS]] {{\[}}[0, 1], [2]]82 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1]]83 // CHECK-NEXT: %[[MATMUL:.+]] = linalg.vecmat 84 // CHECK-SAME: ins(%[[COLLAPSED_LHS]], %[[COLLAPSED_RHS]] : tensor<?xf32>, tensor<?x?xf32>) outs(%[[COLLAPSED_INIT]] : tensor<?xf32>)85 // CHECK-NEXT: %[[DIM1:.*]] = tensor.dim %[[COLLAPSED_INIT]], %[[C0]]86 // CHECK-NEXT: %[[RES:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1]] output_shape [1, %[[DIM1]]]87 // CHECK-NEXT: return %[[RES]]88 %1 = linalg.batch_vecmat ins(%arg0, %arg1 : tensor<1x?xf32>, tensor<1x?x?xf32>)89 outs(%arg2 : tensor<1x?xf32>) -> tensor<1x?xf32>90 return %1 : tensor<1x?xf32>91}92 93// -----94 95func.func @matmul_to_matvec_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x1xf32>, %arg2: tensor<?x1xf32>) -> tensor<?x1xf32> {96 // CHECK-LABEL: @matmul_to_matvec_tensor97 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?xf32>98 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x1xf32>99 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<?x1xf32>100 // CHECK-DAG: %[[C0:.*]] = arith.constant 0101 // CHECK-NEXT: %[[COLLAPSED_RHS:.*]] = tensor.collapse_shape %[[RHS]] {{\[}}[0, 1]]102 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1]]103 // CHECK-NEXT: %[[MATMUL:.+]] = linalg.matvec 104 // CHECK-SAME: ins(%[[LHS]], %[[COLLAPSED_RHS]] : tensor<?x?xf32>, tensor<?xf32>) outs(%[[COLLAPSED_INIT]] : tensor<?xf32>)105 // CHECK-NEXT: %[[DIM0:.*]] = tensor.dim %[[COLLAPSED_INIT]], %[[C0]]106 // CHECK-NEXT: %[[RES:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1]] output_shape [%[[DIM0]], 1]107 // CHECK-NEXT: return %[[RES]]108 %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x1xf32>) outs(%arg2: tensor<?x1xf32>) -> tensor<?x1xf32>109 return %0 : tensor<?x1xf32>110}111 112// -----113 114func.func @matmul_to_matvec(%arg0: memref<?x?xf32>, %arg1: memref<?x1xf32>, %arg2: memref<?x1xf32>) {115 // CHECK-LABEL: @matmul_to_matvec116 // CHECK: linalg.matvec117 linalg.matmul ins(%arg0, %arg1: memref<?x?xf32>, memref<?x1xf32>) outs(%arg2: memref<?x1xf32>)118 return119}120 121// -----122 123func.func @negative_matmul_to_matvec(%arg0: memref<?xf32>, %arg1: memref<?x1xf32>, %arg2: memref<?x1xf32>) {124 // CHECK-LABEL: @negative_matmul_to_matvec125 // CHECK-NOT: linalg.matvec126 linalg.matmul indexing_maps = [127 affine_map<(d0, d1, d2) -> (d2)>,128 affine_map<(d0, d1, d2) -> (d2, d1)>,129 affine_map<(d0, d1, d2) -> (d0, d1)>130 ]131 ins(%arg0, %arg1: memref<?xf32>, memref<?x1xf32>) outs(%arg2: memref<?x1xf32>)132 return133}134 135// -----136 137func.func @matmul_to_vecmat_tensor(%arg0: tensor<1x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<1x?xf32>) -> tensor<1x?xf32> {138 // CHECK-LABEL: @matmul_to_vecmat139 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<1x?xf32>140 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?xf32>141 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<1x?xf32>142 // CHECK-DAG: %[[C0:.*]] = arith.constant 0143 // CHECK-NEXT: %[[COLLAPSED_LHS:.*]] = tensor.collapse_shape %[[LHS]] {{\[}}[0, 1]]144 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1]]145 // CHECK-NEXT: %[[RESULT:.*]] = linalg.vecmat 146 // CHECK-SAME: ins(%[[COLLAPSED_LHS]], %[[RHS]] : tensor<?xf32>, tensor<?x?xf32>) outs(%[[COLLAPSED_INIT]] : tensor<?xf32>)147 // CHECK-NEXT: %[[DIM1:.*]] = tensor.dim %[[COLLAPSED_INIT]], %[[C0]]148 // CHECK-NEXT: %[[RES:.*]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]] output_shape [1, %[[DIM1]]]149 // CHECK-NEXT: return %[[RES]]150 %0 = linalg.matmul ins(%arg0, %arg1: tensor<1x?xf32>, tensor<?x?xf32>) outs(%arg2: tensor<1x?xf32>) -> tensor<1x?xf32>151 return %0 : tensor<1x?xf32>152}153 154// -----155 156func.func @batch_matmul_to_vecmat(%arg0: memref<1x1x?xf32>, %arg1: memref<1x?x?xf32>, %arg2: memref<1x1x?xf32>) {157 // CHECK-LABEL: @batch_matmul_to_vecmat158 // CHECK: linalg.vecmat159 linalg.batch_matmul ins(%arg0, %arg1: memref<1x1x?xf32>, memref<1x?x?xf32>) outs(%arg2: memref<1x1x?xf32>)160 return161}162 163// -----164 165func.func @matvec_to_dot(%arg0: memref<1x?xf32>, %arg1: memref<?xf32>, %arg2: memref<1xf32>) {166 // CHECK-LABEL: @matvec_to_dot167 // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: memref<1x?xf32>168 // CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: memref<?xf32>169 // CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: memref<1xf32>170 // CHECK-NEXT: %[[COLLAPSED_LHS:.*]] = memref.collapse_shape %[[LHS]] {{\[}}[0, 1]]171 // CHECK-NEXT: %[[COLLAPSED_INIT:.*]] = memref.collapse_shape %[[INIT]] []172 // CHECK-NEXT: linalg.dot ins(%[[COLLAPSED_LHS]], %[[RHS]] : memref<?xf32>, memref<?xf32>) outs(%[[COLLAPSED_INIT]] : memref<f32>)173 linalg.matvec ins(%arg0, %arg1: memref<1x?xf32>, memref<?xf32>) outs(%arg2: memref<1xf32>)174 return175}176 177// -----178 179func.func @vecmat_to_dot(%arg0: memref<?xf32>, %arg1: memref<?x1xf32>, %arg2: memref<1xf32>) {180 // CHECK-LABEL: @vecmat_to_dot181 // CHECK: linalg.dot182 linalg.vecmat ins(%arg0, %arg1: memref<?xf32>, memref<?x1xf32>) outs(%arg2: memref<1xf32>)183 return184}185 186// -----187 188func.func @matvec_to_dot_tensor(%arg0: tensor<1x?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<1xf32>) -> tensor<1xf32> {189 // CHECK-LABEL: @matvec_to_dot_tensor190 // CHECK: linalg.dot191 %0 = linalg.matvec ins(%arg0, %arg1: tensor<1x?xf32>, tensor<?xf32>) outs(%arg2: tensor<1xf32>) -> tensor<1xf32>192 return %0 : tensor<1xf32>193}194 195// -----196 197func.func @nonsingleton_batch_matmul(%arg0 : tensor<2x?x?xf32>, %arg1 : tensor<2x?x?xf32>, %arg2: tensor<2x?x?xf32>) -> tensor<2x?x?xf32> {198 // CHECK-LABEL: @nonsingleton_batch_matmul199 // CHECK-NOT: collapse_shape200 // CHECK: linalg.batch_matmul201 // CHECK-NOT: expand_shape202 %1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<2x?x?xf32>, tensor<2x?x?xf32>)203 outs(%arg2 : tensor<2x?x?xf32>) -> tensor<2x?x?xf32>204 return %1 : tensor<2x?x?xf32>205}206 207// -----208 209func.func @nonsingleton_batch_matmul_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>, %arg2: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {210 // CHECK-LABEL: @nonsingleton_batch_matmul_dynamic211 // CHECK-NOT: collapse_shape212 // CHECK: linalg.batch_matmul213 // CHECK-NOT: expand_shape214 %1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)215 outs(%arg2 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>216 return %1 : tensor<?x?x?xf32>217}218