mlir/test/Dialect/Tensor/transform-op-casting.mlir

   1 // RUN: mlir-opt %s --transform-interpreter -allow-unregistered-dialect --split-input-file | FileCheck %s
   2
   3 func.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
   8 func.func private @concat_replacement(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
   9
  10 module 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
  37 func.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
  42 func.func private @collapse_replacement(%arg0: tensor<4x5xf32>) -> tensor<20xf32>
  43
  44 module 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>