105 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s2 3// CHECK-LABEL: @promote_in04// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x42xf32>, %{{.*}}, %{{.*}})5// CHECK: %[[C0:.+]] = arith.constant 06// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]]7// CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor(%[[DIM]]) {memory_space = 1 : i64}8// CHECK: %[[MAT:.+]] = bufferization.materialize_in_destination %[[ARG0]] in %[[ALLOC]]9// CHECK: linalg.matmul ins(%[[MAT]], %{{.*}}10func.func @promote_in0(%arg0: tensor<?x42xf32>, %arg1: tensor<42x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {11 %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x42xf32>, tensor<42x?xf32>)12 outs(%arg2: tensor<?x?xf32>) -> tensor<?x?xf32>13 return %0 : tensor<?x?xf32>14}15 16module attributes {transform.with_named_sequence} {17 transform.named_sequence @__transform_main(%root: !transform.any_op) {18 %mm = transform.structured.match ops{["linalg.matmul"]} in %root19 : (!transform.any_op) -> !transform.any_op20 %op0 = transform.get_operand %mm[0]21 : (!transform.any_op) -> !transform.any_value22 transform.structured.promote_tensor to 1 %op0 : !transform.any_value23 transform.yield24 }25}26 27// -----28 29// CHECK-LABEL: @promote_out30// CHECK-SAME: (%{{.*}}: tensor<?x42xf32>, %{{.*}}: tensor<?x42xf32>, %[[ARG2:.+]]: tensor<?x?xf32>)31func.func @promote_out(%arg0: tensor<?x42xf32>, %arg1: tensor<?x42xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {32 // CHECK: %[[C0:.+]] = arith.constant 033 // CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG2]], %[[C0]]34 // CHECK: %[[C1:.+]] = arith.constant 135 // CHECK: %[[DIM1:.+]] = tensor.dim %[[ARG2]], %[[C1]]36 // CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor(%[[DIM0]], %[[DIM1]]) {memory_space = 1 : i64}37 // CHECK-NOT: materialize_in_destination38 // CHECK: linalg.add {{.*}} outs(%[[ALLOC]]39 %0 = linalg.add ins(%arg0, %arg1 : tensor<?x42xf32>, tensor<?x42xf32>)40 outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>41 return %0 : tensor<?x?xf32>42}43 44module attributes {transform.with_named_sequence} {45 transform.named_sequence @__transform_main(%root: !transform.any_op) {46 %la = transform.structured.match ops{["linalg.add"]} in %root47 : (!transform.any_op) -> !transform.any_op48 %init = transform.get_operand %la[2]49 : (!transform.any_op) -> !transform.any_value50 transform.structured.promote_tensor to 1 %init : !transform.any_value51 52 transform.yield53 }54}55 56// -----57 58// CHECK-LABEL: @promote_in0_out_bufferize59// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x42xf32>, %{{.*}}: tensor<42x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>)60func.func @promote_in0_out_bufferize(%arg0: tensor<?x42xf32>, %arg1: tensor<42x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {61 // CHECK: %[[IN1:.+]] = bufferization.to_buffer %arg1 : tensor<42x?xf32> to memref<42x?xf32, strided<[?, ?], offset: ?>>62 // CHECK: %[[IN0:.+]] = bufferization.to_buffer %arg0 : tensor<?x42xf32> to memref<?x42xf32, strided<[?, ?], offset: ?>>63 // CHECK: %{{.+}} = bufferization.to_buffer %arg0 : tensor<?x42xf32> to memref<?x42xf32, strided<[?, ?], offset: ?>>64 // CHECK: %{{.+}} = bufferization.to_buffer %arg2 : tensor<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>65 // CHECK: %{{.+}} = bufferization.to_buffer %arg2 : tensor<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>66 // CHECK: %[[C0:.+]] = arith.constant 0 : index67 // CHECK: %{{.+}} = memref.dim %{{.+}}, %[[C0]] : memref<?x?xf32, strided<[?, ?], offset: ?>>68 // CHECK: %[[C1:.+]] = arith.constant 1 : index69 // CHECK: %{{.+}} = memref.dim %{{.+}}, %[[C1]] : memref<?x?xf32, strided<[?, ?], offset: ?>>70 // CHECK: %[[ALLOC_OUT:.+]] = memref.alloc(%{{.+}}, %{{.+}}) {alignment = 64 : i64} : memref<?x?xf32, 1>71 // CHECK: %{{.+}} = arith.constant 0 : index72 // CHECK: %{{.+}} = memref.dim %{{.+}}, %{{.+}} : memref<?x42xf32, strided<[?, ?], offset: ?>>73 // CHECK: %[[ALLOC_IN:.+]] = memref.alloc(%{{.+}}) {alignment = 64 : i64} : memref<?x42xf32, 1>74 // CHECK: memref.copy %[[IN0]], %[[ALLOC_IN]] : memref<?x42xf32, strided<[?, ?], offset: ?>> to memref<?x42xf32, 1>75 // CHECK: linalg.add ins(%[[ALLOC_IN]], %[[IN1]] : memref<?x42xf32, 1>, memref<42x?xf32, strided<[?, ?], offset: ?>>) outs(%[[ALLOC_OUT]] : memref<?x?xf32, 1>)76 %0 = linalg.add ins(%arg0, %arg1: tensor<?x42xf32>, tensor<42x?xf32>)77 outs(%arg2: tensor<?x?xf32>) -> tensor<?x?xf32>78 return %0 : tensor<?x?xf32>79}80 81module attributes {transform.with_named_sequence} {82 transform.named_sequence @__transform_main(%root: !transform.any_op) {83 %la = transform.structured.match ops{["linalg.add"]} in %root84 : (!transform.any_op) -> !transform.any_op85 %op0 = transform.get_operand %la[0]86 : (!transform.any_op) -> !transform.any_value87 transform.structured.promote_tensor to 1 %op0 : !transform.any_value88 89 %init = transform.get_operand %la[2]90 : (!transform.any_op) -> !transform.any_value91 transform.structured.promote_tensor to 1 %init : !transform.any_value92 93 %func = transform.structured.match ops{["func.func"]} in %root94 : (!transform.any_op) -> !transform.any_op95 96 %bufferized = transform.bufferization.one_shot_bufferize %func97 : (!transform.any_op) -> !transform.any_op98 99 transform.yield100 }101}102 103 104 105