insert_layout_conversion_output.mlir

#map = affine_map<(d0, d1) -> (d0 * 32 + d1)>
#map1 = affine_map<(d0) -> (d0)>
#map2 = affine_map<(d0) -> (d0 * 32)>
module {
  func.func @insert_conversion(%arg0: !secret.secret<tensor<32x32xi16>> {tensor_ext.layout = #tensor_ext.layout<layout = (d0, d1) -> (d0 * 32 + d1)>}, %arg1: !secret.secret<tensor<32x32xi16>> {tensor_ext.layout = #tensor_ext.layout<layout = (d0, d1) -> (d0 * 32 + d1)>}) -> (!secret.secret<tensor<32xi16>> {tensor_ext.layout = #tensor_ext.layout<layout = (d0) -> (d0)>}) {
    %0 = secret.generic ins(%arg0, %arg1 : !secret.secret<tensor<32x32xi16>>, !secret.secret<tensor<32x32xi16>>) attrs = {arg0 = {layout = #map}, arg1 = {layout = #map}, layout = [#map1]} {
    ^body(%input0: tensor<32x32xi16>, %input1: tensor<32x32xi16>):
      %1 = tensor_ext.sum %input0, 0 {layout = [#map1]} : tensor<32x32xi16> -> tensor<32xi16>
      %2 = tensor_ext.sum %input1, 1 {layout = [#map2]} : tensor<32x32xi16> -> tensor<32xi16>
      %3 = tensor_ext.convert_layout %2 {from_layout = #map2, layout = [#map1], to_layout = #map1} : tensor<32xi16>
      %4 = arith.addi %1, %3 {layout = [#map1]} : tensor<32xi16>
      secret.yield %4 : tensor<32xi16>
    } -> !secret.secret<tensor<32xi16>>
    return %0 : !secret.secret<tensor<32xi16>>
  }
}