// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-named \
// RUN: | FileCheck %s --check-prefix=NAMED,ALL

// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-category \
// RUN: | FileCheck %s --check-prefix=CATEGORY,ALL

#umap = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @unary_ops(%A: tensor<?x?x?xf32>, %Out: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#umap, #umap],
    iterator_types = ["parallel", "parallel","parallel"]}
    ins(%A : tensor<?x?x?xf32>)
    outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.exp %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %1 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%0 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.log %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %2 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%1 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.absf %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %3 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%2 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.ceil %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %4 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%3 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.floor %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %5 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%4 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = arith.negf %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %cst_1 = arith.constant 1.0 : f32
  %6 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%5 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = arith.divf %cst_1, %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %7 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%6 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.round %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %8 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%7 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.sqrt %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %9 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%8 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.rsqrt %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %10 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%9 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = arith.mulf %in, %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %11 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%10 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.tanh %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  %12 = linalg.generic
          {indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
          ins(%11 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
  ^bb0(%in: f32, %out: f32):
    %v = math.erf %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?x?xf32>
  return %12 : tensor<?x?x?xf32>
}

// ALL-LABEL: unary_ops
// ALL-SAME: %[[A:.+]]: tensor<?x?x?xf32>, %[[OUT:.+]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: %[[RES0:.+]] = linalg.exp
// NAMED-SAME: ins(%[[A]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES1:.+]] = linalg.log
// NAMED-SAME: ins(%[[RES0]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES2:.+]] = linalg.abs
// NAMED-SAME: ins(%[[RES1]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES3:.+]] = linalg.ceil
// NAMED-SAME: ins(%[[RES2]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES4:.+]] = linalg.floor
// NAMED-SAME: ins(%[[RES3]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES5:.+]] = linalg.negf
// NAMED-SAME: ins(%[[RES4]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES6:.+]] = linalg.reciprocal
// NAMED-SAME: ins(%[[RES5]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES7:.+]] = linalg.round
// NAMED-SAME: ins(%[[RES6]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES8:.+]] = linalg.sqrt
// NAMED-SAME: ins(%[[RES7]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES9:.+]] = linalg.rsqrt
// NAMED-SAME: ins(%[[RES8]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES10:.+]] = linalg.square
// NAMED-SAME: ins(%[[RES9]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES11:.+]] = linalg.tanh
// NAMED-SAME: ins(%[[RES10]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// NAMED: %[[RES12:.+]] = linalg.erf
// NAMED-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<exp>
// CATEGORY-SAME: ins(%[[A]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<log>
// CATEGORY-SAME: ins(%[[RES0]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<abs>
// CATEGORY-SAME: ins(%[[RES1]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<ceil>
// CATEGORY-SAME: ins(%[[RES2]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES4:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<floor>
// CATEGORY-SAME: ins(%[[RES3]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES5:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<negf>
// CATEGORY-SAME: ins(%[[RES4]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES6:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<reciprocal>
// CATEGORY-SAME: ins(%[[RES5]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES7:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<round>
// CATEGORY-SAME: ins(%[[RES6]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES8:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sqrt>
// CATEGORY-SAME: ins(%[[RES7]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES9:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<rsqrt>
// CATEGORY-SAME: ins(%[[RES8]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES10:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<square>
// CATEGORY-SAME: ins(%[[RES9]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES11:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<tanh>
// CATEGORY-SAME: ins(%[[RES10]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
// CATEGORY: %[[RES12:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<erf>
// CATEGORY-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>

// -----

func.func @unary_ops_non_identity(%A: tensor<?xf32>, %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>],
    iterator_types = ["parallel", "parallel"]}
    ins(%A : tensor<?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %out: f32):  
    %v = math.exp %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}


// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0, d1) -> (d1)>
// ALL-DAG: #[[MAP_TP:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// ALL: unary_ops_non_identity
// ALL-SAME: %[[A:.+]]: tensor<?xf32>, %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>

// Named ops cannot carry user-defined indexing maps -> expect no change.
// NAMED-NOT: linalg.exp
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.elementwise kind=#linalg.elementwise_kind<exp>
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]]]
// CATEGORY-SAME: ins(%[[A]] : tensor<?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// -----

#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @binary_ops_int(%A: tensor<?x?xi32>, %B: tensor<?x?xi32>,
                          %Out: tensor<?x?xi32>) -> tensor<?x?xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.addi %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %1 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%0, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.subi %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %2 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%1, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.muli %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %3 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%2, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.divsi %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %4 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%3, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.divui %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %5 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%4, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.maxsi %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %6 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%5, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.minsi %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  return %6 : tensor<?x?xi32>
}

// ALL-LABEL: binary_ops_int
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi32>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// NAMED-NOT: linalg.generic
// NAMED: %[[RES0:.+]] = linalg.add
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES1:.+]] = linalg.sub
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES2:.+]] = linalg.mul
// NAMED-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES3:.+]] = linalg.div
// NAMED-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES4:.+]] = linalg.div_unsigned
// NAMED-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES5:.+]] = linalg.max
// NAMED-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES6:.+]] = linalg.min
// NAMED-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div>
// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES4:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div_unsigned>
// CATEGORY-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES5:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
// CATEGORY-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES6:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
// CATEGORY-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @binary_ops_float(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                            %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.addf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %1 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%0, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.subf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %2 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%1, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.mulf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %3 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%2, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.divf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %4 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%3, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.maximumf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %5 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%4, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.minimumf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %6 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%5, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = math.powf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  return %6 : tensor<?x?xf32>
}

// ALL-LABEL: binary_ops_float
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// NAMED-NOT: linalg.generic
// NAMED: %[[RES0:.+]] = linalg.add
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES1:.+]] = linalg.sub
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES2:.+]] = linalg.mul
// NAMED-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES3:.+]] = linalg.div
// NAMED-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES4:.+]] = linalg.max
// NAMED-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES5:.+]] = linalg.min
// NAMED-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES6:.+]] = linalg.powf
// NAMED-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div>
// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES4:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
// CATEGORY-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES5:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
// CATEGORY-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES6:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<powf>
// CATEGORY-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @binary_ops_complex(%A: tensor<?x?xcomplex<f32>>,
                              %B: tensor<?x?xcomplex<f32>>,
                              %Out: tensor<?x?xcomplex<f32>>)
                               -> tensor<?x?xcomplex<f32>> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
    outs(%Out : tensor<?x?xcomplex<f32>>) {
  ^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
    %v = complex.add %in, %in_0 : complex<f32>
    linalg.yield %v : complex<f32>
  } -> tensor<?x?xcomplex<f32>>
  %1 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%0, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
    outs(%Out : tensor<?x?xcomplex<f32>>) {
  ^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
    %v = complex.sub %in, %in_0 : complex<f32>
    linalg.yield %v : complex<f32>
  } -> tensor<?x?xcomplex<f32>>
  %2 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%1, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
    outs(%Out : tensor<?x?xcomplex<f32>>) {
  ^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
    %v = complex.mul %in, %in_0 : complex<f32>
    linalg.yield %v : complex<f32>
  } -> tensor<?x?xcomplex<f32>>
  %3 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%2, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
    outs(%Out : tensor<?x?xcomplex<f32>>) {
  ^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
    %v = complex.div %in, %in_0 : complex<f32>
    linalg.yield %v : complex<f32>
  } -> tensor<?x?xcomplex<f32>>
  return %3 : tensor<?x?xcomplex<f32>>
}

// ALL-LABEL: binary_ops_complex
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xcomplex<f32>>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// NAMED-NOT: linalg.generic
// NAMED: %[[RES0:.+]] = linalg.add
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES1:.+]] = linalg.sub
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES2:.+]] = linalg.mul
// NAMED-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES3:.+]] = linalg.div
// NAMED-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div>
// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @binary_ops_bool(%A: tensor<?x?xi1>, %B: tensor<?x?xi1>,
                           %Out: tensor<?x?xi1>) -> tensor<?x?xi1> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xi1>, tensor<?x?xi1>)
    outs(%Out : tensor<?x?xi1>) {
  ^bb0(%in: i1, %in_0: i1, %out: i1):
    %v = arith.ori %in, %in_0 : i1
    linalg.yield %v : i1
  } -> tensor<?x?xi1>
  %1 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%0, %B : tensor<?x?xi1>, tensor<?x?xi1>)
    outs(%Out : tensor<?x?xi1>) {
  ^bb0(%in: i1, %in_0: i1, %out: i1):
    %v = arith.andi %in, %in_0 : i1
    linalg.yield %v : i1
  } -> tensor<?x?xi1>
  return %1 : tensor<?x?xi1>
}

// ALL-LABEL: binary_ops_bool
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi1>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// NAMED-NOT: linalg.generic
// NAMED: %[[RES0:.+]] = linalg.add
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED: %[[RES1:.+]] = linalg.mul
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @binary_ops_uint(%A: tensor<?x?xi32>, %B: tensor<?x?xi32>,
                          %Out: tensor<?x?xi32>) -> tensor<?x?xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.maxui %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  %1 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%0, %B : tensor<?x?xi32>, tensor<?x?xi32>)
    outs(%Out : tensor<?x?xi32>) {
  ^bb0(%in: i32, %in_0: i32, %out: i32):
    %v = arith.minui %in, %in_0 : i32
    linalg.yield %v : i32
  } -> tensor<?x?xi32>
  return %1 : tensor<?x?xi32>
}

// ALL-LABEL: binary_ops_uint
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi32>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// No named ops yet for unsigned max/min -> expect no change.
// NAMED-NOT: linalg.{{max|min}}
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<max_unsigned>
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<min_unsigned>
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

func.func @binary_ops_non_identity(%A: tensor<?xf32>, %B: tensor<?x?xf32>,
                                   %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>,
                      affine_map<(d0, d1) -> (d0, d1)>],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):  
    %v = arith.addf %in, %in_0 : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0, d1) -> (d1)>
// ALL-DAG: #[[MAP_TP:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// ALL-DAG: #[[MAP_ID:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// ALL: binary_ops_non_identity
// ALL-SAME: %[[A:.+]]: [[TTY1D:tensor<\?xf32>]], %[[B:.+]]: [[TTY:tensor<\?x\?xf32>]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// Named ops cannot carry user-defined indexing maps -> expect no change.
// NAMED-NOT: linalg.add
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]], #[[MAP_ID]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY1D]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0, d1) -> (d0, d1)>
#bcast = affine_map<(d0, d1) -> (d0)>
func.func @binary_ops_swapped(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                              %C: tensor<?xf32>, %Out: tensor<?x?xf32>)
                               -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.addf %in_0, %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  %1 = linalg.generic
    {indexing_maps = [#map, #bcast, #map],
    iterator_types = ["parallel", "parallel"]}
    ins(%0, %C : tensor<?x?xf32>, tensor<?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %v = arith.subf %in_0, %in : f32
    linalg.yield %v : f32
  } -> tensor<?x?xf32>
  return %1 : tensor<?x?xf32>
}

// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0, d1) -> (d0)>
// ALL-DAG: #[[MAP_ID:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// ALL: binary_ops_swapped
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[B:.+]]: [[TTY]],
// ALL-SAME: %[[C:.+]]: [[TTY1D:tensor<\?xf32>]],
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// NAMED: %[[RES0:.+]] = linalg.add
// NAMED-SAME: ins(%[[B]], %[[A]] : [[TTY]], [[TTY]])
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// NAMED-NOT: linalg.sub
// NAMED: linalg.generic
// NAMED-SAME: ins(%[[RES0]], %[[C]] : [[TTY]], [[TTY1D]])

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: ins(%[[B]], %[[A]] : [[TTY]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
// CATEGORY-SAME: ins(%[[C]], %[[RES0]] : [[TTY1D]], [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0) -> (d0)>
func.func @unary_op_with_scalar(%A: tensor<?xi32>, %Out: tensor<?xi32>)
                                 -> tensor<?xi32> {
  %cst = arith.constant 123 : i32
  %0 = linalg.generic
    {indexing_maps = [#map, #map],
    iterator_types = ["parallel"]}
    ins(%A : tensor<?xi32>)
    outs(%Out : tensor<?xi32>) {
  ^bb0(%in: i32, %out: i32):
    %v = arith.addi %cst, %in : i32
    linalg.yield %v : i32
  } -> tensor<?xi32>
  return %0 : tensor<?xi32>
}

// CATEGORY-DAG: #[[MAP_ID:.+]] = affine_map<(d0) -> (d0)>
// CATEGORY-DAG: #[[MAP_BC:.+]] = affine_map<(d0) -> ()>
// ALL: unary_op_with_scalar
// CATEGORY-SAME: %[[A:.+]]: [[TTY:tensor<\?xi32>]],
// CATEGORY-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]

// Named ops cannot broadcast from a scalar operand -> expect no change.
// NAMED-NOT: linalg.add
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[CST:.+]] = arith.constant 123 : i32
// CATEGORY: linalg.elementwise kind=#linalg.elementwise_kind<add>
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
// CATEGORY-SAME: ins(%[[CST]], %[[A]] : i32, [[TTY]])
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]

// -----

#map = affine_map<(d0) -> (d0)>
func.func @negative_unary_op_using_block_arg_twice(%A: tensor<?xi32>,
                                                   %Out: tensor<?xi32>)
                                                    -> tensor<?xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map],
    iterator_types = ["parallel"]}
    ins(%A : tensor<?xi32>)
    outs(%Out : tensor<?xi32>) {
  ^bb0(%in: i32, %out: i32):
    %v = arith.addi %in, %in : i32
    linalg.yield %v : i32
  } -> tensor<?xi32>
  return %0 : tensor<?xi32>
}

// ALL-LABEL: negative_unary_op_using_block_arg_twice

// Named ops cannot broadcast from a scalar operand -> expect no change.
// NAMED-NOT: linalg.add

// There is no scalar operand to hoist -> expect no change.
// CATEGORY-NOT: linalg.elementwise kind=#linalg.elementwise_kind<add>

// ALL: linalg.generic

// -----

///----------------------------------------------------------------------------------------
/// Tests for linalg.matmul
///----------------------------------------------------------------------------------------

#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @op_matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                     %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>

// ALL-LABEL: op_matmul
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>,
// ALL-SAME: %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// Cast-auditing tests: ensure we only specialize when the cast semantics can
// be expressed by linalg.matmul, and use the cast attribute when needed.

// Check matmul with unsigned cast is correctly raised back to named op.
func.func @op_matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi32>,
                                   %Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi32>)
    outs(%Out : tensor<16x32xi32>) {
  ^bb0(%in: i16, %in_0: i32, %out: i32):
    %1 = arith.extui %in : i16 to i32
    %3 = arith.muli %1, %in_0 : i32
    %4 = arith.addi %out, %3 : i32
    linalg.yield %4 : i32
  } -> tensor<16x32xi32>
  return %0 : tensor<16x32xi32>
}

// ALL-LABEL: op_matmul_unsigned_cast

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
// CATEGORY-SAME: {cast = #linalg.type_fn<cast_unsigned>}

// Ensures truncation rounding is tolerated with unsigned cases.
// Note: We only consider casts as conflicting if they have different
// signedness behaviours, and then we do not specialize if they do
// conflict. Since this is not such a case, we do not block specialization.
// Also the roundtrip lowering back to linalg.generic for such an op is
// expected to produce the same thing again, so we are not loosing
// information here.
func.func @op_matmul_unsigned_cast_and_truncate(%A: tensor<16x8xi16>, %B: tensor<8x32xi64>,
                                                %Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi64>)
    outs(%Out : tensor<16x32xi32>) {
  ^bb0(%in: i16, %in_0: i64, %out: i32):
    %1 = arith.extui %in : i16 to i32
    %2 = arith.trunci %in_0 : i64 to i32
    %3 = arith.muli %1, %2 : i32
    %4 = arith.addi %out, %3 : i32
    linalg.yield %4 : i32
  } -> tensor<16x32xi32>
  return %0 : tensor<16x32xi32>
}

// ALL-LABEL: op_matmul_unsigned_cast_and_truncate

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
// CATEGORY-SAME: {cast = #linalg.type_fn<cast_unsigned>}

// Signed casts are the default, no cast attribute is required.
func.func @op_matmul_signed_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
                                 %Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
    outs(%Out : tensor<16x32xi32>) {
  ^bb0(%in: i16, %in_0: i16, %out: i32):
    %1 = arith.extsi %in : i16 to i32
    %2 = arith.extsi %in_0 : i16 to i32
    %3 = arith.muli %1, %2 : i32
    %4 = arith.addi %out, %3 : i32
    linalg.yield %4 : i32
  } -> tensor<16x32xi32>
  return %0 : tensor<16x32xi32>
}

// ALL-LABEL: op_matmul_signed_cast

// NAMED-NOT: linalg.generic
// NAMED-NOT: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
// NAMED: linalg.matmul

// CATEGORY-NOT: linalg.generic
// CATEGORY: %[[RES:.+]] = linalg.contract
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
// CATEGORY-NOT: {cast =
// CATEGORY-SAME: ins
// CATEGORY: return %[[RES]]

// Mixed signed/unsigned inputs cannot be encoded with a single cast attribute.
func.func @negative_op_matmul_mixed_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
                                         %Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
    outs(%Out : tensor<16x32xi32>) {
  ^bb0(%in: i16, %in_0: i16, %out: i32):
    %1 = arith.extui %in : i16 to i32
    %2 = arith.extsi %in_0 : i16 to i32
    %3 = arith.muli %1, %2 : i32
    %4 = arith.addi %out, %3 : i32
    linalg.yield %4 : i32
  } -> tensor<16x32xi32>
  return %0 : tensor<16x32xi32>
}

// ALL-LABEL: negative_op_matmul_mixed_cast

// NAMED: linalg.generic
// NAMED-NOT: linalg.matmul

// CATEGORY: linalg.generic
// CATEGORY-NOT: linalg.contract

// Output-side casts are not representable by the named matmul ops.
func.func @negative_op_matmul_output_cast(%A: tensor<16x8xi32>, %B: tensor<8x32xi32>,
                                          %Out: tensor<16x32xi64>) -> tensor<16x32xi64> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi32>, tensor<8x32xi32>)
    outs(%Out : tensor<16x32xi64>) {
  ^bb0(%in: i32, %in_0: i32, %out: i64):
    %3 = arith.trunci %out : i64 to i32
    %4 = arith.muli %in, %in_0 : i32
    %5 = arith.addi %3, %4 : i32
    %6 = arith.extsi %5 : i32 to i64
    linalg.yield %6 : i64
  } -> tensor<16x32xi64>
  return %0 : tensor<16x32xi64>
}

// ALL-LABEL: negative_op_matmul_output_cast

// NAMED: linalg.generic
// NAMED-NOT: linalg.matmul

// CATEGORY: linalg.generic
// CATEGORY-NOT: linalg.contract

// Bitcasts are not modeled by the cast attribute, but should not block
// specialization.
// NOTE: Bitcasts are not preserved by the matmul named op during
// roundtrip, so this is potentially loosing information here.
// See #177593 for more details.
func.func @op_matmul_bitcast_int_to_float(%A: tensor<16x8xi32>,
                                          %B: tensor<8x32xi32>,
                                          %Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi32>, tensor<8x32xi32>)
    outs(%Out : tensor<16x32xf32>) {
  ^bb0(%in: i32, %in_0: i32, %out: f32):
    %1 = arith.bitcast %in : i32 to f32
    %2 = arith.bitcast %in_0 : i32 to f32
    %3 = arith.mulf %1, %2 : f32
    %4 = arith.addf %out, %3 : f32
    linalg.yield %4 : f32
  } -> tensor<16x32xf32>
  return %0 : tensor<16x32xf32>
}

// ALL-LABEL: op_matmul_bitcast_int_to_float

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// Signed float casts only use sitofp, which defaults to signed semantics.
func.func @op_matmul_signed_cast_float(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
                                       %Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
    outs(%Out : tensor<16x32xf32>) {
  ^bb0(%in: i16, %in_0: i16, %out: f32):
    %1 = arith.sitofp %in : i16 to f32
    %2 = arith.sitofp %in_0 : i16 to f32
    %3 = arith.mulf %1, %2 : f32
    %4 = arith.addf %out, %3 : f32
    linalg.yield %4 : f32
  } -> tensor<16x32xf32>
  return %0 : tensor<16x32xf32>
}

// ALL-LABEL: op_matmul_signed_cast_float

// NAMED-NOT: linalg.generic
// NAMED-NOT: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
// NAMED: linalg.matmul

// CATEGORY-NOT: linalg.generic
// CATEGORY-NOT: linalg.contract{{.*}}{cast =
// CATEGORY: %[[RES:.+]] = linalg.contract
// CATEGORY: return %[[RES]]

// Unsigned float casts are expressed via uitofp and use the unsigned cast attr.
func.func @op_matmul_unsigned_cast_float(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
                                         %Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
    outs(%Out : tensor<16x32xf32>) {
  ^bb0(%in: i16, %in_0: i16, %out: f32):
    %1 = arith.uitofp %in : i16 to f32
    %2 = arith.uitofp %in_0 : i16 to f32
    %3 = arith.mulf %1, %2 : f32
    %4 = arith.addf %out, %3 : f32
    linalg.yield %4 : f32
  } -> tensor<16x32xf32>
  return %0 : tensor<16x32xf32>
}

// ALL-LABEL: op_matmul_unsigned_cast_float

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract{{.*}}{cast = #linalg.type_fn<cast_unsigned>}

// -----

///----------------------------------------------------------------------------------------
/// Tests for linalg.batch_matmul
///----------------------------------------------------------------------------------------

#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
func.func @op_batch_matmul(%A: tensor<2x16x8xf32>, %B: tensor<2x8x16xf32>,
                           %Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
    outs(%Out : tensor<2x16x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x16x16xf32>
  return %0 : tensor<2x16x16xf32>
}

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>

// ALL-LABEL: op_batch_matmul
// ALL-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>,  %[[OUT:.+]]: tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.batch_matmul ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>) outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// Ensure that the unsigned cast path for cast detection is exercised for
// batch_matmul as well.
func.func @op_batch_matmul_unsigned_cast(%A: tensor<2x16x8xi16>,
                                         %B: tensor<2x8x16xi64>,
                                         %Out: tensor<2x16x16xi32>) -> tensor<2x16x16xi32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<2x16x8xi16>, tensor<2x8x16xi64>)
    outs(%Out : tensor<2x16x16xi32>) {
  ^bb0(%in: i16, %in_0: i64, %out: i32):
    %1 = arith.extui %in : i16 to i32
    %2 = arith.trunci %in_0 : i64 to i32
    %3 = arith.muli %1, %2 : i32
    %4 = arith.addi %out, %3 : i32
    linalg.yield %4 : i32
  } -> tensor<2x16x16xi32>
  return %0 : tensor<2x16x16xi32>
}

// ALL-LABEL: op_batch_matmul_unsigned_cast

// NAMED-NOT: linalg.generic
// NAMED: linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
// CATEGORY-SAME: {cast = #linalg.type_fn<cast_unsigned>}

// -----

// A multi-reduction contraction.
#mapA = affine_map<(m, n, k1, k2) -> (m, k1, k2)>
#mapB = affine_map<(m, n, k1, k2) -> (k2, k1, n)>
#mapC = affine_map<(m, n, k1, k2) -> (m, n)>
func.func @op_multi_reduction(%A: tensor<10x20x30xf32>,
                              %B: tensor<30x20x40xf32>,
                              %C: tensor<10x40xf32>) -> tensor<10x40xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapA, #mapB, #mapC],
    iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
    ins(%A, %B : tensor<10x20x30xf32>, tensor<30x20x40xf32>)
    outs(%C : tensor<10x40xf32>) {
  ^bb0(%a: f32, %b: f32, %c: f32):
    %1 = arith.mulf %a, %b : f32
    %2 = arith.addf %c, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<10x40xf32>
  return %0 : tensor<10x40xf32>
}

// ALL-LABEL: op_multi_reduction

// Cannot be lifted to named matrix multiply.
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// -----

// Batch dim not in identity position: batch dim d0 appears at result
// position 1 in A's map instead of position 0.
#mapBni0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
#mapBni1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#mapBni2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
func.func @batch_matmul_non_identity_batch(%A: tensor<4x2x8xf32>, %B: tensor<2x8x16xf32>,
                                           %Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapBni0, #mapBni1, #mapBni2],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<4x2x8xf32>, tensor<2x8x16xf32>)
    outs(%Out : tensor<2x4x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x4x16xf32>
  return %0 : tensor<2x4x16xf32>
}

// ALL-LABEL: batch_matmul_non_identity_batch

// Cannot be lifted to named matrix multiply.
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// -----

// TODO: named matvec
#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d1)>
#map2 = affine_map<(d0, d1) -> (d0)>
func.func @op_matvec(%A: tensor<?x?xf32>, %B: tensor<?xf32>, %Out: tensor<?xf32>) -> tensor<?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map, #map1, #map2],
    iterator_types = ["parallel", "reduction"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?xf32>)
    outs(%Out : tensor<?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?xf32>
  return %0 : tensor<?xf32>
}

// ALL-LABEL: op_matvec

// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// -----

// Matmul transpose A: A is accessed as (k, m) instead of (m, k)
#map_ta = affine_map<(d0, d1, d2) -> (d2, d0)>
#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @op_matmul_transpose_a(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                 %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_ta, #map_b, #map_c],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// ALL-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
// ALL-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>

// ALL-LABEL: op_matmul_transpose_a
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]

// -----

// Matmul transpose B: B is accessed as (n, k) instead of (k, n)
#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
#map_tb = affine_map<(d0, d1, d2) -> (d1, d2)>
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @op_matmul_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                 %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_a, #map_tb, #map_c],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// ALL-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// ALL-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>

// ALL-LABEL: op_matmul_transpose_b
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// -----

// Batch matmul transpose A: A is accessed as (b, k, m) instead of (b, m, k)
#map_ta = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
#map_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
func.func @op_batch_matmul_transpose_a(%A: tensor<2x8x4xf32>, %B: tensor<2x8x16xf32>,
                                       %Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_ta, #map_b, #map_c],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<2x8x4xf32>, tensor<2x8x16xf32>)
    outs(%Out : tensor<2x4x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x4x16xf32>
  return %0 : tensor<2x4x16xf32>
}

// ALL-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
// ALL-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>

// ALL-LABEL: op_batch_matmul_transpose_a
// ALL-SAME: %[[A:.+]]: tensor<2x8x4xf32>, %[[B:.+]]: tensor<2x8x16xf32>, %[[OUT:.+]]: tensor<2x4x16xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.batch_matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x8x4xf32>, tensor<2x8x16xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x8x4xf32>, tensor<2x8x16xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>

// -----

// Batch matmul transpose B: B is accessed as (b, n, k) instead of (b, k, n)
#map_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map_tb = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
#map_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
func.func @op_batch_matmul_transpose_b(%A: tensor<2x4x8xf32>, %B: tensor<2x16x8xf32>,
                                       %Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_a, #map_tb, #map_c],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<2x4x8xf32>, tensor<2x16x8xf32>)
    outs(%Out : tensor<2x4x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x4x16xf32>
  return %0 : tensor<2x4x16xf32>
}

// ALL-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// ALL-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>

// ALL-LABEL: op_batch_matmul_transpose_b
// ALL-SAME: %[[A:.+]]: tensor<2x4x8xf32>, %[[B:.+]]: tensor<2x16x8xf32>, %[[OUT:.+]]: tensor<2x4x16xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.batch_matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x16x8xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x16x8xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>

// -----

// Both A and B transposed.
#map_ta = affine_map<(d0, d1, d2) -> (d2, d0)>
#map_tb = affine_map<(d0, d1, d2) -> (d1, d2)>
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @op_matmul_transpose_a_and_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                       %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_ta, #map_tb, #map_c],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// ALL-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
// ALL-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>

// ALL-LABEL: op_matmul_transpose_a_and_b
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_TB]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_TB]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// -----

// Output transposed: C is accessed as (n, m) instead of (m, n).
#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
#map_tc = affine_map<(d0, d1, d2) -> (d1, d0)>
func.func @op_matmul_transposed_output(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                       %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_a, #map_b, #map_tc],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// ALL-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// ALL-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// ALL-DAG: #[[$MAP_TC:.+]] = affine_map<(d0, d1, d2) -> (d1, d0)>

// ALL-LABEL: op_matmul_transposed_output
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_TC]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_TC]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// -----

// Matmul with non-canonical loop ordering.
#map_nc_a = affine_map<(m, k, n) -> (m, k)>
#map_nc_b = affine_map<(m, k, n) -> (k, n)>
#map_nc_c = affine_map<(m, k, n) -> (m, n)>
func.func @op_matmul_non_canonical_loops(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                         %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_nc_a, #map_nc_b, #map_nc_c],
    iterator_types = ["parallel", "reduction", "parallel"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>

// ALL-LABEL: op_matmul_non_canonical_loops
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// -----

// Batch matmul with non-canonical loop ordering.
#map_bnc_a = affine_map<(batch, m, k, n) -> (batch, m, k)>
#map_bnc_b = affine_map<(batch, m, k, n) -> (batch, k, n)>
#map_bnc_c = affine_map<(batch, m, k, n) -> (batch, m, n)>
func.func @op_batch_matmul_non_canonical_loops(%A: tensor<2x16x8xf32>, %B: tensor<2x8x16xf32>,
                                               %Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_bnc_a, #map_bnc_b, #map_bnc_c],
    iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
    ins(%A, %B : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
    outs(%Out : tensor<2x16x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x16x16xf32>
  return %0 : tensor<2x16x16xf32>
}

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>

// ALL-LABEL: op_batch_matmul_non_canonical_loops
// ALL-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>, %[[OUT:.+]]: tensor<2x16x16xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.batch_matmul
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// -----

// Matmul with non-canonical loop ordering (d0=m, d1=k, d2=n) and B transposed.
#map_nc_tb_a = affine_map<(d0, d1, d2) -> (d0, d1)>
#map_nc_tb_b = affine_map<(d0, d1, d2) -> (d2, d1)>
#map_nc_tb_c = affine_map<(d0, d1, d2) -> (d0, d2)>
func.func @op_matmul_non_canonical_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                               %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_nc_tb_a, #map_nc_tb_b, #map_nc_tb_c],
    iterator_types = ["parallel", "reduction", "parallel"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// NAMED-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// NAMED-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
// NAMED-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CATEGORY-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>

// ALL-LABEL: op_matmul_non_canonical_transpose_b
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// -----

// Batch matmul with non-canonical loop ordering (d0=batch, d1=m, d2=k, d3=n)
// and B Transposed.
#map_bnc_tb_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
#map_bnc_tb_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map_bnc_tb_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
func.func @op_batch_matmul_non_canonical_transpose_b(%A: tensor<2x16x8xf32>, %B: tensor<2x16x8xf32>,
                                                     %Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_bnc_tb_a, #map_bnc_tb_b, #map_bnc_tb_c],
    iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
    ins(%A, %B : tensor<2x16x8xf32>, tensor<2x16x8xf32>)
    outs(%Out : tensor<2x16x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x16x16xf32>
  return %0 : tensor<2x16x16xf32>
}

// NAMED-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// NAMED-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// NAMED-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CATEGORY-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>

// ALL-LABEL: op_batch_matmul_non_canonical_transpose_b
// ALL-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x16x8xf32>, %[[OUT:.+]]: tensor<2x16x16xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.batch_matmul
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x16x8xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x16x8xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>

// -----

// Matmul with fully permuted loop ordering.
#map_fs_a = affine_map<(d0, d1, d2) -> (d1, d0)>
#map_fs_b = affine_map<(d0, d1, d2) -> (d0, d2)>
#map_fs_c = affine_map<(d0, d1, d2) -> (d1, d2)>
func.func @op_matmul_fully_shuffled_loops(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
                                          %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_fs_a, #map_fs_b, #map_fs_c],
    iterator_types = ["reduction", "parallel", "parallel"]}
    ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d1, d0)>
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>

// ALL-LABEL: op_matmul_fully_shuffled_loops
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>

// NAMED-NOT: linalg.generic
// NAMED: linalg.matmul
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>

// -----

// TODO: this could also be specialized to a named matmul.
#map_bcast_a = affine_map<(d0, d1, d2) -> (d2)>
#map_bcast_b = affine_map<(d0, d1, d2) -> (d2, d1)>
#map_bcast_c = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @op_matmul_broadcast_a(%A: tensor<?xf32>, %B: tensor<?x?xf32>,
                                 %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_bcast_a, #map_bcast_b, #map_bcast_c],
    iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?xf32>, tensor<?x?xf32>)
    outs(%Out : tensor<?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// ALL-LABEL: op_matmul_broadcast_a

// NAMED: linalg.generic
// NAMED-NOT: linalg.matmul

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// -----

// TODO: this could also be specialized to a named batch_matmul.
#map_bbcast_a = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
#map_bbcast_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map_bbcast_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
func.func @op_batch_matmul_broadcast_a(%A: tensor<16x8xf32>, %B: tensor<2x8x16xf32>,
                                       %Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_bbcast_a, #map_bbcast_b, #map_bbcast_c],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<16x8xf32>, tensor<2x8x16xf32>)
    outs(%Out : tensor<2x16x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x16x16xf32>
  return %0 : tensor<2x16x16xf32>
}

// ALL-LABEL: op_batch_matmul_broadcast_a

// NAMED: linalg.generic
// NAMED-NOT: linalg.batch_matmul

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// -----

// TODO: this could also be specialized to a named batch_matmul.
#map_bbcast2_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map_bbcast2_b = affine_map<(d0, d1, d2, d3) -> (d3)>
#map_bbcast2_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
func.func @op_batch_matmul_broadcast_b(%A: tensor<2x16x8xf32>, %B: tensor<8xf32>,
                                       %Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_bbcast2_a, #map_bbcast2_b, #map_bbcast2_c],
    iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<2x16x8xf32>, tensor<8xf32>)
    outs(%Out : tensor<2x16x16xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<2x16x16xf32>
  return %0 : tensor<2x16x16xf32>
}

// ALL-LABEL: op_batch_matmul_broadcast_b

// NAMED: linalg.generic
// NAMED-NOT: linalg.batch_matmul

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// -----

///----------------------------------------------------------------------------------------
/// Tests for linalg.mmt4d
///----------------------------------------------------------------------------------------

#mapA = affine_map<(m, n, k, m0, n0, k0) -> (m, k, m0, k0)>
#mapB = affine_map<(m, n, k, m0, n0, k0) -> (n, k, n0, k0)>
#mapC = affine_map<(m, n, k, m0, n0, k0) -> (m, n, m0, n0)>
func.func @op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
                    %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapA, #mapB, #mapC],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: op_mmt4d

// NAMED-NOT: linalg.generic
// NAMED: linalg.mmt4d

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// MMT4D transpose A inner and outer:
//   A is accessed as (k, m, k0, m0) instead of (m, k, m0, k0)
#map_tA = affine_map<(m, n, k, m0, n0, k0) -> (k, m, k0, m0)>
func.func @op_mmt4d_transpose_a(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
                                %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_tA, #mapB, #mapC],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: op_mmt4d_transpose_a

// NAMED-NOT: linalg.generic
// NAMED: linalg.mmt4d

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// MMT4D transpose B inner and outer:
//   B is accessed as (k, n, k0, n0) instead of (n, k, n0, k0)
#map_tB = affine_map<(m, n, k, m0, n0, k0) -> (k, n, k0, n0)>
func.func @op_mmt4d_transpose_b(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
                                %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapA, #map_tB, #mapC],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: op_mmt4d_transpose_b

// NAMED-NOT: linalg.generic
// NAMED: linalg.mmt4d

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// MMT4D transpose both A and B inner and outer:
func.func @op_mmt4d_transpose_a_and_b(
    %A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
    %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#map_tA, #map_tB, #mapC],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: op_mmt4d_transpose_a_and_b

// NAMED-NOT: linalg.generic
// NAMED: linalg.mmt4d

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// MMT4D transpose C inner and outer:
//   C is accessed as (n, m, n0, m0) instead of (m, n, m0, n0)
#map_tC = affine_map<(m, n, k, m0, n0, k0) -> (n, m, n0, m0)>
func.func @op_mmt4d_transpose_c(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
                                %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapA, #mapB, #map_tC],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: op_mmt4d_transpose_c

// NAMED-NOT: linalg.generic
// NAMED: linalg.mmt4d

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// MMT4D transpose C inner only:
//   C is accessed as (m, n, n0, m0) instead of (m, n, m0, n0)
#map_tC_inner = affine_map<(m, n, k, m0, n0, k0) -> (m, n, n0, m0)>
func.func @op_mmt4d_transpose_c_inner(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
                                      %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapA, #mapB, #map_tC_inner],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: op_mmt4d_transpose_c_inner

// NAMED-NOT: linalg.generic
// NAMED: linalg.mmt4d

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract

// Negative MMT4D:
// A can only be accessed as inner transpose or outer transpose of (m, k, m0, k0)
#mapA_negative = affine_map<(m, n, k, m0, n0, k0) -> (n, k, n0, k0)>
func.func @negative_op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
                             %C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
  %0 = linalg.generic
    {indexing_maps = [#mapA_negative, #mapB, #mapC],
    iterator_types = ["parallel", "parallel", "reduction",
                      "parallel", "parallel", "reduction"]}
    ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
    outs(%C : tensor<?x?x?x?xf32>) {
  ^bb0(%in: f32, %in_0: f32, %out: f32):
    %1 = arith.mulf %in, %in_0 : f32
    %2 = arith.addf %out, %1 : f32
    linalg.yield %2 : f32
  } -> tensor<?x?x?x?xf32>
  return %0 : tensor<?x?x?x?xf32>
}

// ALL-LABEL: negative_op_mmt4d

// NAMED-NOT: linalg.mmt4d
// NAMED: linalg.generic

// CATEGORY-NOT: linalg.generic
// CATEGORY: linalg.contract