1// RUN: mlir-opt %s --transform-interpreter -allow-unregistered-dialect --split-input-file | FileCheck %s 2 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_op 13 %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_op 15 %ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value 16 %out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value 17 transform.func.cast_and_call %f#1(%ins) -> %out before %concat { 18 transform.type_conversion.tensor.cast_shape_dynamic_dims 19 } : (!transform.any_op, !transform.any_value, 20 !transform.any_value, !transform.any_op) -> !transform.any_op 21 transform.apply_dce to %f#0 : !transform.any_op 22 transform.yield 23 } 24} 25 26// CHECK-LABEL: func.func @cast_to_dynamic 27// 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_op 47 %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_op 49 %ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value 50 %out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value 51 transform.func.cast_and_call %f#1(%ins) -> %out before %concat { 52 transform.type_conversion.tensor.cast_shape_dynamic_dims ignore_dynamic_info 53 } : (!transform.any_op, !transform.any_value, 54 !transform.any_value, !transform.any_op) -> !transform.any_op 55 transform.apply_dce to %f#0 : !transform.any_op 56 transform.yield 57 } 58} 59 60// CHECK-LABEL: func.func @cast_to_static 61// 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