Nicolas Vasilache | 11fb8d0 | 2022-03-21 20:33:30 +0100 | [diff] [blame^] | 1 | // RUN: iree-dialects-opt -linalg-interp-transforms %s | FileCheck %s |
| 2 | |
| 3 | // CHECK-LABEL: func @matmul_tensors |
| 4 | // CHECK-NOT: linalg |
| 5 | // CHECK: llvm |
| 6 | func @matmul_tensors( |
| 7 | %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32> { linalg.inplaceable = true}) |
| 8 | -> tensor<128x128xf32> { |
| 9 | %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) |
| 10 | outs(%arg2: tensor<128x128xf32>) |
| 11 | -> tensor<128x128xf32> |
| 12 | |
| 13 | return %0 : tensor<128x128xf32> |
| 14 | } |
| 15 | |
| 16 | |
| 17 | pdl.pattern @pdl_target : benefit(1) { |
| 18 | %args = operands |
| 19 | %results = types |
| 20 | %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>) |
| 21 | apply_native_constraint "nestedInFunc"[@matmul_tensors](%0 : !pdl.operation) |
| 22 | // TODO: we don't want this, but it is the required terminator for pdl.pattern |
| 23 | rewrite %0 with "iree_linalg_transform.apply" |
| 24 | } |
| 25 | |
| 26 | iree_linalg_transform.sequence { |
| 27 | %0 = match @pdl_target |
| 28 | %1 = tile %0 {sizes = [4, 4, 4]} |
| 29 | %2 = vectorize %1 {vectorize_padding = true} |
| 30 | bufferize |
| 31 | lower_vectors { multireduction_lowering = "innerreduce"} |
| 32 | lower_to_llvm |
| 33 | } |