// RUN: mlir-opt -test-linalg-drop-unit-dims --split-input-file %s | FileCheck %s --check-prefixes=CHECK,NOENCODE
// RUN: mlir-opt -test-linalg-drop-unit-dims=collapse-encoded --split-input-file %s | FileCheck %s --check-prefixes=CHECK,ENCODE

// Drop only the outermost unit dimension (controlled using a control function)
// This test does not use an encoding, therefore behavior in both modes is identical.
func.func @drop_outermost_unit_dims(%arg0: tensor<1x1x42xf32>) -> tensor<1x1x42xf32> {
  %0 = tensor.empty() : tensor<1x1x42xf32>
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    ins(%arg0 : tensor<1x1x42xf32>) outs(%0 : tensor<1x1x42xf32>) {
      ^bb0(%b0: f32, %b1 : f32):
        %2 = arith.addf %b0, %b1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32>
  return %1 : tensor<1x1x42xf32>
}
// CHECK-LABEL: func @drop_outermost_unit_dims
//  CHECK-SAME:     %[[ARG0:.+]]: tensor<1x1x42xf32>
//       CHECK:   %[[OUTS:.+]] = tensor.empty()
//       CHECK:   %[[ARG0_RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2]{{\]}}
//       CHECK:   %[[OUTS_RESHAPE:.+]] = tensor.collapse_shape %[[OUTS]] {{\[}}[0, 1], [2]{{\]}}
//       CHECK:   %[[GENERIC:.+]] = linalg.generic
//  CHECK-SAME:       ins(%[[ARG0_RESHAPE]] :
//  CHECK-SAME:       outs(%[[OUTS_RESHAPE]] :
//       CHECK:   %[[EXPAND_SHAPE:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0, 1], [2]{{\]}}
//       CHECK:   return %[[EXPAND_SHAPE]]

// -----

// Drop outermost unit dimension with operand that has an encoding.
// With the default behavior, the transformation is aborted and operation remains unchanged.
// With the custom behavior, the operand gets collapsed and encoding is preserved in the collapse.

#encoding = #test.tensor_encoding<"encoding">

func.func @drop_unit_dims_encoded_operand(%arg0: tensor<1x1x42xf32>, %arg1: tensor<1x1x42xf32, #encoding>) -> tensor<1x1x42xf32> {
  %0 = tensor.empty() : tensor<1x1x42xf32>
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    ins(%arg0, %arg1 : tensor<1x1x42xf32>, tensor<1x1x42xf32, #encoding>) outs(%0 : tensor<1x1x42xf32>) {
      ^bb0(%in0: f32, %in1 : f32, %out : f32):
        %2 = arith.addf %in0, %in1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32>
  return %1 : tensor<1x1x42xf32>
}

// NOENCODE-LABEL:   @drop_unit_dims_encoded_operand(
// NOENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32>,
// NOENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) -> tensor<1x1x42xf32> {
// NOENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32>
// NOENCODE:           %[[GENERIC_0:.*]] = linalg.generic
// NOENCODE-SAME:        iterator_types = ["parallel", "parallel", "parallel"]}
// NOENCODE-SAME:        ins(%[[ARG0]], %[[ARG1]] : tensor<1x1x42xf32>, tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>)
// NOENCODE-SAME:        outs(%[[EMPTY_0]] : tensor<1x1x42xf32>)
// NOENCODE:           return %[[GENERIC_0]] : tensor<1x1x42xf32>

//   ENCODE-LABEL:   @drop_unit_dims_encoded_operand(
//   ENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32>,
//   ENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>) -> tensor<1x1x42xf32> {
//   ENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32>
//   ENCODE:           %[[COLLAPSE_SHAPE_0:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
//   ENCODE:           %[[COLLAPSE_SHAPE_1:.*]] = tensor.collapse_shape %[[ARG1]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">> into tensor<1x42xf32, #test.tensor_encoding<"encoding">>
//   ENCODE:           %[[COLLAPSE_SHAPE_2:.*]] = tensor.collapse_shape %[[EMPTY_0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
//   ENCODE:           %[[GENERIC_0:.*]] = linalg.generic
//   ENCODE-SAME:        iterator_types = ["parallel", "parallel"]
//   ENCODE-SAME:        ins(%[[COLLAPSE_SHAPE_0]], %[[COLLAPSE_SHAPE_1]] : tensor<1x42xf32>, tensor<1x42xf32, #test.tensor_encoding<"encoding">>)
//   ENCODE-SAME:        outs(%[[COLLAPSE_SHAPE_2]] : tensor<1x42xf32>)
//   ENCODE:           %[[EXPAND_SHAPE_0:.*]] = tensor.expand_shape %[[GENERIC_0]] {{\[\[}}0, 1], [2]] output_shape [1, 1, 42] : tensor<1x42xf32> into tensor<1x1x42xf32>
//   ENCODE:           return %[[EXPAND_SHAPE_0]] : tensor<1x1x42xf32>

// -----

// Drop outermost unit dimension with result that has an encoding.
// With the default behavior, the transformation is aborted and operation remains unchanged.
// With the custom behavior, the result gets expanded and encoding is preserved in the expansion.

#encoding = #test.tensor_encoding<"encoding">

func.func @drop_unit_dims_encoded_result(%arg0: tensor<1x1x42xf32>, %arg1: tensor<1x1x42xf32>) -> tensor<1x1x42xf32, #encoding> {
  %0 = tensor.empty() : tensor<1x1x42xf32, #encoding>
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    ins(%arg0, %arg1 : tensor<1x1x42xf32>, tensor<1x1x42xf32>) outs(%0 : tensor<1x1x42xf32, #encoding>) {
      ^bb0(%in0: f32, %in1 : f32, %out : f32):
        %2 = arith.addf %in0, %in1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32, #encoding>
  return %1 : tensor<1x1x42xf32, #encoding>
}

// NOENCODE-LABEL:   @drop_unit_dims_encoded_result(
// NOENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32>,
// NOENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32>) -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
// NOENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
// NOENCODE:           %[[GENERIC_0:.*]] = linalg.generic
// NOENCODE-SAME:        iterator_types = ["parallel", "parallel", "parallel"]
// NOENCODE-SAME:        ins(%[[ARG0]], %[[ARG1]] : tensor<1x1x42xf32>, tensor<1x1x42xf32>)
// NOENCODE-SAME:        outs(%[[EMPTY_0]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>)
// NOENCODE-NOT:       tensor.expand_shape
// NOENCODE:           return %[[GENERIC_0]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>

//   ENCODE-LABEL:   @drop_unit_dims_encoded_result(
//   ENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32>,
//   ENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32>) -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
//   ENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
//   ENCODE:           %[[COLLAPSE_SHAPE_0:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
//   ENCODE:           %[[COLLAPSE_SHAPE_1:.*]] = tensor.collapse_shape %[[ARG1]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
//   ENCODE:           %[[COLLAPSE_SHAPE_2:.*]] = tensor.collapse_shape %[[EMPTY_0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">> into tensor<1x42xf32, #test.tensor_encoding<"encoding">>
//   ENCODE:           %[[GENERIC_0:.*]] = linalg.generic
//   ENCODE-SAME:        iterator_types = ["parallel", "parallel"]
//   ENCODE-SAME:        ins(%[[COLLAPSE_SHAPE_0]], %[[COLLAPSE_SHAPE_1]] : tensor<1x42xf32>, tensor<1x42xf32>)
//   ENCODE-SAME:        outs(%[[COLLAPSE_SHAPE_2]] : tensor<1x42xf32, #test.tensor_encoding<"encoding">>)
//   ENCODE:           %[[EXPAND_SHAPE_0:.*]] = tensor.expand_shape %[[GENERIC_0]] {{\[\[}}0, 1], [2]] output_shape [1, 1, 42] : tensor<1x42xf32, #test.tensor_encoding<"encoding">> into tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>
//   ENCODE:           return %[[EXPAND_SHAPE_0]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding">>

// -----

#encoding1 = #test.tensor_encoding<"encoding1">
#encoding2 = #test.tensor_encoding<"encoding2">

func.func @drop_unit_dims_encoded_operand_and_result(%arg0: tensor<1x1x42xf32>, %arg1: tensor<1x1x42xf32, #encoding1>) -> tensor<1x1x42xf32, #encoding2> {
  %0 = tensor.empty() : tensor<1x1x42xf32, #encoding2>
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    ins(%arg0, %arg1 : tensor<1x1x42xf32>, tensor<1x1x42xf32, #encoding1>) outs(%0 : tensor<1x1x42xf32, #encoding2>) {
      ^bb0(%in0: f32, %in1 : f32, %out : f32):
        %2 = arith.addf %in0, %in1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32, #encoding2>
  return %1 : tensor<1x1x42xf32, #encoding2>
}

// NOENCODE-LABEL:   @drop_unit_dims_encoded_operand_and_result
// NOENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32>
// NOENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>
// NOENCODE-SAME:      -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
// NOENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
// NOENCODE:           %[[GENERIC_0:.*]] = linalg.generic
// NOENCODE-SAME:        iterator_types = ["parallel", "parallel", "parallel"]
// NOENCODE-SAME:        ins(%[[ARG0]], %[[ARG1]] : tensor<1x1x42xf32>, tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>)
// NOENCODE-SAME:        outs(%[[EMPTY_0]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>)
// NOENCODE-NOT:       tensor.expand_shape
// NOENCODE:           return %[[GENERIC_0]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>

//   ENCODE-LABEL:   @drop_unit_dims_encoded_operand_and_result
//   ENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32>,
//   ENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>
//   ENCODE-SAME:      -> tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE:           %[[COLLAPSE_SHAPE_0:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
//   ENCODE:           %[[COLLAPSE_SHAPE_1:.*]] = tensor.collapse_shape %[[ARG1]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>
//   ENCODE-SAME:        into tensor<1x42xf32, #test.tensor_encoding<"encoding1">>
//   ENCODE:           %[[COLLAPSE_SHAPE_2:.*]] = tensor.collapse_shape %[[EMPTY_0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE-SAME:        into tensor<1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE:           %[[GENERIC_0:.*]] = linalg.generic
//   ENCODE-SAME:        iterator_types = ["parallel", "parallel"]
//   ENCODE-SAME:        ins(%[[COLLAPSE_SHAPE_0]], %[[COLLAPSE_SHAPE_1]] : tensor<1x42xf32>, tensor<1x42xf32, #test.tensor_encoding<"encoding1">>)
//   ENCODE-SAME:        outs(%[[COLLAPSE_SHAPE_2]] : tensor<1x42xf32, #test.tensor_encoding<"encoding2">>)
//   ENCODE:           %[[EXPAND_SHAPE_0:.*]] = tensor.expand_shape %[[GENERIC_0]] {{\[\[}}0, 1], [2]] output_shape [1, 1, 42] : tensor<1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE-SAME:        into tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE:           return %[[EXPAND_SHAPE_0]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>

// -----

#encoding1 = #test.tensor_encoding<"encoding1">
#encoding2 = #test.tensor_encoding<"encoding2">

func.func @drop_unit_dims_two_encoded_operands(%arg0: tensor<1x1x42xf32, #encoding1>, %arg1: tensor<1x1x42xf32, #encoding2>) -> tensor<1x1x42xf32> {
  %0 = tensor.empty() : tensor<1x1x42xf32>
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    ins(%arg0, %arg1 : tensor<1x1x42xf32, #encoding1>, tensor<1x1x42xf32, #encoding2>) outs(%0 : tensor<1x1x42xf32>) {
      ^bb0(%in0: f32, %in1 : f32, %out : f32):
        %2 = arith.addf %in0, %in1 : f32
        linalg.yield %2 : f32
    } -> tensor<1x1x42xf32>
  return %1 : tensor<1x1x42xf32>
}

// NOENCODE-LABEL:   @drop_unit_dims_two_encoded_operands
// NOENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>
// NOENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>) -> tensor<1x1x42xf32>
// NOENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32>
// NOENCODE:           %[[GENERIC_0:.*]] = linalg.generic
// NOENCODE-SAME:        iterator_types = ["parallel", "parallel", "parallel"]}
// NOENCODE-SAME:        ins(%[[ARG0]], %[[ARG1]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>, tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>)
// NOENCODE-SAME:        outs(%[[EMPTY_0]] : tensor<1x1x42xf32>)
// NOENCODE:           return %[[GENERIC_0]] : tensor<1x1x42xf32>

//   ENCODE-LABEL:   @drop_unit_dims_two_encoded_operands
//   ENCODE-SAME:      %[[ARG0:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>
//   ENCODE-SAME:      %[[ARG1:.*]]: tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>) -> tensor<1x1x42xf32>
//   ENCODE:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<1x1x42xf32>
//   ENCODE:           %[[COLLAPSE_SHAPE_0:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding1">>
//   ENCODE-SAME:        into tensor<1x42xf32, #test.tensor_encoding<"encoding1">>
//   ENCODE:           %[[COLLAPSE_SHAPE_1:.*]] = tensor.collapse_shape %[[ARG1]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE-SAME:        into tensor<1x42xf32, #test.tensor_encoding<"encoding2">>
//   ENCODE:           %[[COLLAPSE_SHAPE_2:.*]] = tensor.collapse_shape %[[EMPTY_0]] {{\[\[}}0, 1], [2]] : tensor<1x1x42xf32> into tensor<1x42xf32>
//   ENCODE:           %[[GENERIC_0:.*]] = linalg.generic
//   ENCODE-SAME:        iterator_types = ["parallel", "parallel"]
//   ENCODE-SAME:        ins(%[[COLLAPSE_SHAPE_0]], %[[COLLAPSE_SHAPE_1]] : tensor<1x42xf32, #test.tensor_encoding<"encoding1">>, tensor<1x42xf32, #test.tensor_encoding<"encoding2">>)
//   ENCODE-SAME:        outs(%[[COLLAPSE_SHAPE_2]] : tensor<1x42xf32>)
//   ENCODE:           %[[EXPAND_SHAPE_0:.*]] = tensor.expand_shape %[[GENERIC_0]] {{\[\[}}0, 1], [2]] output_shape [1, 1, 42] : tensor<1x42xf32> into tensor<1x1x42xf32>
//   ENCODE:           return %[[EXPAND_SHAPE_0]] : tensor<1x1x42xf32>