// RUN: mlir-opt %s --transform-interpreter -allow-unregistered-dialect --split-input-file | FileCheck %s

func.func @cast_to_dynamic(%arg0: tensor<10x13xf32>, %arg1: tensor<3x13xf32>) -> tensor<13x13xf32> {
  %0 = tensor.concat dim(0) %arg0, %arg1 : (tensor<10x13xf32>, tensor<3x13xf32>) -> tensor<13x13xf32>
  func.return %0 : tensor<13x13xf32>
}

func.func private @concat_replacement(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32>

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
    %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
    %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    %concat = transform.structured.match ops{["tensor.concat"]} in %f#0 : (!transform.any_op) -> !transform.any_op
    %ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value
    %out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value
    transform.func.cast_and_call %f#1(%ins) -> %out before %concat {
      transform.type_conversion.tensor.cast_shape_dynamic_dims
    } : (!transform.any_op, !transform.any_value,
         !transform.any_value, !transform.any_op) -> !transform.any_op
    transform.apply_dce to %f#0 : !transform.any_op
    transform.yield
  }
}

// CHECK-LABEL: func.func @cast_to_dynamic
//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<10x13xf32>
//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9]+]]: tensor<3x13xf32>
//   CHECK-DAG:   %[[CAST0:.+]] = tensor.cast %[[ARG0]] : tensor<10x13xf32> to tensor<?x?xf32>
//   CHECK-DAG:   %[[CAST1:.+]] = tensor.cast %[[ARG1]] : tensor<3x13xf32> to tensor<?x?xf32>
//       CHECK:   %[[CALL:.+]] = call @concat_replacement(%[[CAST0]], %[[CAST1]])
//       CHECK:   %[[CAST_RES:.+]] = tensor.cast %[[CALL]] : tensor<?x?xf32> to tensor<13x13xf32>
//       CHECK:   return %[[CAST_RES]] : tensor<13x13xf32>

// -----

func.func @cast_to_static(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {
  %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
  func.return %0 : tensor<?xf32>
}

func.func private @collapse_replacement(%arg0: tensor<4x5xf32>) -> tensor<20xf32>

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
    %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
    %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    %concat = transform.structured.match ops{["tensor.collapse_shape"]} in %f#0 : (!transform.any_op) -> !transform.any_op
    %ins = transform.get_operand %concat[all] : (!transform.any_op) -> !transform.any_value
    %out = transform.get_result %concat[all] : (!transform.any_op) -> !transform.any_value
    transform.func.cast_and_call %f#1(%ins) -> %out before %concat {
      transform.type_conversion.tensor.cast_shape_dynamic_dims ignore_dynamic_info
    } : (!transform.any_op, !transform.any_value,
         !transform.any_value, !transform.any_op) -> !transform.any_op
    transform.apply_dce to %f#0 : !transform.any_op
    transform.yield
  }
}

// CHECK-LABEL: func.func @cast_to_static
//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
//   CHECK-DAG:   %[[CAST_IN:.+]] = tensor.cast %[[ARG0]] : tensor<?x?xf32> to tensor<4x5xf32>
//       CHECK:   %[[CALL:.+]] = call @collapse_replacement(%[[CAST_IN]])
//       CHECK:   %[[CAST_RES:.+]] = tensor.cast %[[CALL]] : tensor<20xf32> to tensor<?xf32>
//       CHECK:   return %[[CAST_RES]] : tensor<?xf32>