//===----------------------------------------------------------------------===//
// Utility Methods
//===----------------------------------------------------------------------===//

func.func private @generate_2D_source_f16(%height : index, %width : index) -> tensor<?x?xf16> {
  %init_source = tensor.empty(%height, %width) : tensor<?x?xf16>
  %source = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]}
      outs(%init_source : tensor<?x?xf16>) {
    ^bb0(%b0 : f16):
      %outer = linalg.index 0 : index
      %inner = linalg.index 1 : index
      %strided = arith.muli %outer, %width : index
      %linearized = arith.addi %inner, %strided : index
      %linearized_i16 = arith.index_cast %linearized : index to i16
      %linearized_f16 = arith.sitofp %linearized_i16 : i16 to f16
      linalg.yield %linearized_f16 : f16
  } -> tensor<?x?xf16>
  // This blocks the fusion for inputs and testing ops.
  %0 = util.optimization_barrier %source : tensor<?x?xf16>
  %1 = flow.tensor.tie_shape %0 : tensor<?x?xf16>{%height, %width}
  return %1 : tensor<?x?xf16>
}

func.func private @generate_2D_source_f32(%height : index, %width : index) -> tensor<?x?xf32> {
  %init_source = tensor.empty(%height, %width) : tensor<?x?xf32>
  %source = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]}
      outs(%init_source : tensor<?x?xf32>) {
    ^bb0(%b0 : f32):
      %outer = linalg.index 0 : index
      %inner = linalg.index 1 : index
      %strided = arith.muli %outer, %width : index
      %linearized = arith.addi %inner, %strided : index
      %linearized_i32 = arith.index_cast %linearized : index to i32
      %linearized_f32 = arith.sitofp %linearized_i32 : i32 to f32
      linalg.yield %linearized_f32 : f32
  } -> tensor<?x?xf32>
  // This blocks the fusion for inputs and testing ops.
  %0 = util.optimization_barrier %source : tensor<?x?xf32>
  %1 = flow.tensor.tie_shape %0 : tensor<?x?xf32>{%height, %width}
  return %1 : tensor<?x?xf32>
}

func.func private @generate_2D_source_i8(%height : index, %width : index) -> tensor<?x?xi8> {
  %c255 = arith.constant 255 : index
  %init_source = tensor.empty(%height, %width) : tensor<?x?xi8>
  %source = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]}
      outs(%init_source : tensor<?x?xi8>) {
    ^bb0(%b0 : i8):
      %outer = linalg.index 0 : index
      %inner = linalg.index 1 : index
      %strided = arith.muli %outer, %width : index
      %linearized = arith.addi %inner, %strided : index
      %linearized_rem = arith.remsi %linearized, %c255 : index
      %linearized_i8 = arith.index_cast %linearized_rem : index to i8
      linalg.yield %linearized_i8 : i8
  } -> tensor<?x?xi8>
  // This blocks the fusion for inputs and testing ops.
  %0 = util.optimization_barrier %source : tensor<?x?xi8>
  %1 = flow.tensor.tie_shape %0 : tensor<?x?xi8>{%height, %width}
  return %1 : tensor<?x?xi8>
}

func.func private @generate_2D_source_i32(%height : index, %width : index) -> tensor<?x?xi32> {
  %init_source = tensor.empty(%height, %width) : tensor<?x?xi32>
  %source = linalg.generic {
      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
      iterator_types = ["parallel", "parallel"]}
      outs(%init_source : tensor<?x?xi32>) {
    ^bb0(%b0 : i32):
      %outer = linalg.index 0 : index
      %inner = linalg.index 1 : index
      %strided = arith.muli %outer, %width : index
      %linearized = arith.addi %inner, %strided : index
      %linearized_i32 = arith.index_cast %linearized : index to i32
      linalg.yield %linearized_i32 : i32
  } -> tensor<?x?xi32>
  // This blocks the fusion for inputs and testing ops.
  %0 = util.optimization_barrier %source : tensor<?x?xi32>
  %1 = flow.tensor.tie_shape %0 : tensor<?x?xi32>{%height, %width}
  return %1 : tensor<?x?xi32>
}

//===----------------------------------------------------------------------===//
// f32.f32.f32 variants
//===----------------------------------------------------------------------===//

#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
#encoding_f32f32f32_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f32f32f32_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f32f32f32_acc = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>

func.func @set_encoding_f32f32f32_lhs() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
  %source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f32f32f32_lhs>
  %barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f32f32f32_lhs>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f32f32f32_lhs> -> tensor<129x255xf32>
  check.expect_almost_eq(%2, %source) : tensor<129x255xf32>
  return
}

func.func @set_encoding_f32f32f32_rhs() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
  %source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f32f32f32_rhs>
  %barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f32f32f32_rhs>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f32f32f32_rhs> -> tensor<129x255xf32>
  check.expect_almost_eq(%2, %source) : tensor<129x255xf32>
  return
}

func.func @set_encoding_f32f32f32_acc() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
  %source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f32f32f32_acc>
  %barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f32f32f32_acc>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f32f32f32_acc> -> tensor<129x255xf32>
  check.expect_almost_eq(%2, %source) : tensor<129x255xf32>
  return
}

//===----------------------------------------------------------------------===//
// i8.i8.i32 variants
//===----------------------------------------------------------------------===//

#encoding_i8i8i32_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_i8i8i32_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_i8i8i32_acc = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>

func.func @set_encoding_i8i8i32_lhs() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_i8(%height, %width) : (index, index) -> tensor<?x?xi8>
  %source = tensor.cast %0 : tensor<?x?xi8> to tensor<129x255xi8>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xi8> -> tensor<129x255xi8, #encoding_i8i8i32_lhs>
  %barrire = util.optimization_barrier %1 : tensor<129x255xi8, #encoding_i8i8i32_lhs>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xi8, #encoding_i8i8i32_lhs> -> tensor<129x255xi8>
  check.expect_eq(%2, %source) : tensor<129x255xi8>
  return
}

func.func @set_encoding_i8i8i32_rhs() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_i8(%height, %width) : (index, index) -> tensor<?x?xi8>
  %source = tensor.cast %0 : tensor<?x?xi8> to tensor<129x255xi8>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xi8> -> tensor<129x255xi8, #encoding_i8i8i32_rhs>
  %barrire = util.optimization_barrier %1 : tensor<129x255xi8, #encoding_i8i8i32_rhs>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xi8, #encoding_i8i8i32_rhs> -> tensor<129x255xi8>
  check.expect_eq(%2, %source) : tensor<129x255xi8>
  return
}

func.func @set_encoding_i8i8i32_acc() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_i32(%height, %width) : (index, index) -> tensor<?x?xi32>
  %source = tensor.cast %0 : tensor<?x?xi32> to tensor<129x255xi32>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xi32> -> tensor<129x255xi32, #encoding_i8i8i32_acc>
  %barrire = util.optimization_barrier %1 : tensor<129x255xi32, #encoding_i8i8i32_acc>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xi32, #encoding_i8i8i32_acc> -> tensor<129x255xi32>
  check.expect_eq(%2, %source) : tensor<129x255xi32>
  return
}


//===----------------------------------------------------------------------===//
// f16.f16.f32 variants
//===----------------------------------------------------------------------===//

#encoding_f16f16f32_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f16f16f32_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f16f16f32_acc = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>

func.func @set_encoding_f16f16f32_lhs() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_f16(%height, %width) : (index, index) -> tensor<?x?xf16>
  %source = tensor.cast %0 : tensor<?x?xf16> to tensor<129x255xf16>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xf16> -> tensor<129x255xf16, #encoding_f16f16f32_lhs>
  %barrire = util.optimization_barrier %1 : tensor<129x255xf16, #encoding_f16f16f32_lhs>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xf16, #encoding_f16f16f32_lhs> -> tensor<129x255xf16>
  check.expect_eq(%2, %source) : tensor<129x255xf16>
  return
}

func.func @set_encoding_f16f16f32_rhs() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_f16(%height, %width) : (index, index) -> tensor<?x?xf16>
  %source = tensor.cast %0 : tensor<?x?xf16> to tensor<129x255xf16>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xf16> -> tensor<129x255xf16, #encoding_f16f16f32_rhs>
  %barrire = util.optimization_barrier %1 : tensor<129x255xf16, #encoding_f16f16f32_rhs>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xf16, #encoding_f16f16f32_rhs> -> tensor<129x255xf16>
  check.expect_eq(%2, %source) : tensor<129x255xf16>
  return
}

func.func @set_encoding_f16f16f32_acc() {
  %height = arith.constant 129 : index
  %width = arith.constant 255 : index
  %0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
  %source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>

  %1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f16f16f32_acc>
  %barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f16f16f32_acc>
  %2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f16f16f32_acc> -> tensor<129x255xf32>
  check.expect_eq(%2, %source) : tensor<129x255xf32>
  return
}