66 lines · plain
1// RUN: mlir-opt %s --transform-interpreter -allow-unregistered-dialect --split-input-file | FileCheck %s2 3func.func @cast_to_dynamic(%arg0: tensor<10x13xf32>, %arg1: tensor<3x13xf32>) -> tensor<13x13xf32> {4 %0 = tensor.concat dim(0) %arg0, %arg1 : (tensor<10x13xf32>, tensor<3x13xf32>) -> tensor<13x13xf32>5 func.return %0 : tensor<13x13xf32>6}7 8func.func private @concat_replacement(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32>9 10module attributes {transform.with_named_sequence} {11 transform.named_sequence @__transform_main(%arg0: !transform.any_op) {12 %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op13 %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)14 %concat = transform.structured.match ops{["tensor.concat"]} in %f#0 : (!transform.any_op) -> !transform.any_op15 %ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value16 %out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value17 transform.func.cast_and_call %f#1(%ins) -> %out before %concat {18 transform.type_conversion.tensor.cast_shape_dynamic_dims19 } : (!transform.any_op, !transform.any_value,20 !transform.any_value, !transform.any_op) -> !transform.any_op21 transform.apply_dce to %f#0 : !transform.any_op22 transform.yield23 }24}25 26// CHECK-LABEL: func.func @cast_to_dynamic27// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<10x13xf32>28// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<3x13xf32>29// CHECK-DAG: %[[CAST0:.+]] = tensor.cast %[[ARG0]] : tensor<10x13xf32> to tensor<?x?xf32>30// CHECK-DAG: %[[CAST1:.+]] = tensor.cast %[[ARG1]] : tensor<3x13xf32> to tensor<?x?xf32>31// CHECK: %[[CALL:.+]] = call @concat_replacement(%[[CAST0]], %[[CAST1]])32// CHECK: %[[CAST_RES:.+]] = tensor.cast %[[CALL]] : tensor<?x?xf32> to tensor<13x13xf32>33// CHECK: return %[[CAST_RES]] : tensor<13x13xf32>34 35// -----36 37func.func @cast_to_static(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {38 %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>39 func.return %0 : tensor<?xf32>40}41 42func.func private @collapse_replacement(%arg0: tensor<4x5xf32>) -> tensor<20xf32>43 44module attributes {transform.with_named_sequence} {45 transform.named_sequence @__transform_main(%arg0: !transform.any_op) {46 %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op47 %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)48 %concat = transform.structured.match ops{["tensor.collapse_shape"]} in %f#0 : (!transform.any_op) -> !transform.any_op49 %ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value50 %out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value51 transform.func.cast_and_call %f#1(%ins) -> %out before %concat {52 transform.type_conversion.tensor.cast_shape_dynamic_dims ignore_dynamic_info53 } : (!transform.any_op, !transform.any_value,54 !transform.any_value, !transform.any_op) -> !transform.any_op55 transform.apply_dce to %f#0 : !transform.any_op56 transform.yield57 }58}59 60// CHECK-LABEL: func.func @cast_to_static61// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>62// CHECK-DAG: %[[CAST_IN:.+]] = tensor.cast %[[ARG0]] : tensor<?x?xf32> to tensor<4x5xf32>63// CHECK: %[[CALL:.+]] = call @collapse_replacement(%[[CAST_IN]])64// CHECK: %[[CAST_RES:.+]] = tensor.cast %[[CALL]] : tensor<20xf32> to tensor<?xf32>65// CHECK: return %[[CAST_RES]] : tensor<?xf32>66