Skip to content

Instantly share code, notes, and snippets.

@pashu123

pashu123/new.mlir

Created January 10, 2023 14:32
Show Gist options
  • Save pashu123/e9f40745bf461844c213affa691000fc to your computer and use it in GitHub Desktop.
Save pashu123/e9f40745bf461844c213affa691000fc to your computer and use it in GitHub Desktop.
Download ZIP
Raw
new.mlir
#map = affine_map<(d0, d1, d2, d3) -> (d1)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, 0, 0)>
#map4 = affine_map<(d0, d1, d2, d3) -> (0, d1, 0, 0)>
#map5 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map6 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
#map7 = affine_map<(d0, d1) -> (d0, d1)>
#map8 = affine_map<(d0, d1) -> (d1, d0)>
#map9 = affine_map<(d0, d1, d2) -> (0, d1, d2)>
#map10 = affine_map<(d0, d1, d2) -> (d1, d2)>
#map11 = affine_map<(d0, d1, d2) -> (d2)>
#map12 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>
#map13 = affine_map<(d0, d1, d2) -> (d0, d1, 0)>
#map14 = affine_map<(d0, d1, d2) -> (0, d1, 0)>
#map15 = affine_map<(d0, d1) -> (d1)>
module attributes {torch.debug_module_name = "_lambda"} {
ml_program.global private mutable @global_seed(dense<0> : tensor<i64>) : tensor<i64>
func.func @forward(%arg0: tensor<1x4x128x128xf16>) -> tensor<1x3x512x512xf16> {
%cst = arith.constant dense<[1.343990e-01, 8.300780e-02, 4.443360e-02]> : tensor<3xf16>
%cst_0 = arith.constant dense_resource<__elided__> : tensor<3x128x3x3xf16>
%cst_1 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_2 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_3 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_4 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
%cst_5 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_6 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_7 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_8 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
%cst_9 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_10 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_11 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_12 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
%cst_13 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_14 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_15 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_16 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
%cst_17 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_18 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_19 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_20 = arith.constant dense_resource<__elided__> : tensor<128x256x1x1xf16>
%cst_21 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_22 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
%cst_23 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_24 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_25 = arith.constant dense_resource<__elided__> : tensor<128xf16>
%cst_26 = arith.constant dense_resource<__elided__> : tensor<128x256x3x3xf16>
%cst_27 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_28 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_29 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_30 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
%cst_31 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_32 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
%cst_33 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_34 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_35 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_36 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
%cst_37 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_38 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_39 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_40 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
%cst_41 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_42 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_43 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_44 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
%cst_45 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_46 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_47 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_48 = arith.constant dense_resource<__elided__> : tensor<256x512x1x1xf16>
%cst_49 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_50 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
%cst_51 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_52 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_53 = arith.constant dense_resource<__elided__> : tensor<256xf16>
%cst_54 = arith.constant dense_resource<__elided__> : tensor<256x512x3x3xf16>
%cst_55 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_56 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_57 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_58 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_59 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_60 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_61 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_62 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_63 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_64 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_65 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_66 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_67 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_68 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_69 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_70 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_71 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_72 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_73 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_74 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_75 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_76 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_77 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_78 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_79 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_80 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_81 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_82 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_83 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_84 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_85 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_86 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_87 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_88 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_89 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_90 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_91 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_92 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
%cst_93 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_94 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
%cst_95 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_96 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
%cst_97 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_98 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
%cst_99 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_100 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_101 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_102 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_103 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_104 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_105 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_106 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
%cst_107 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_108 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_109 = arith.constant dense_resource<__elided__> : tensor<512xf16>
%cst_110 = arith.constant dense_resource<__elided__> : tensor<512x4x3x3xf16>
%cst_111 = arith.constant dense<[2.230220e-01, 4.936520e-01, 1.726070e-01, -1.182860e-01]> : tensor<4xf16>
%cst_112 = arith.constant dense_resource<__elided__> : tensor<4x4x1x1xf16>
%cst_113 = arith.constant 0.000000e+00 : f16
%cst_114 = arith.constant 0.000000e+00 : f64
%cst_115 = arith.constant 1.000000e+00 : f16
%cst_116 = arith.constant -6.550400e+04 : f16
%cst_117 = arith.constant 9.9999999999999995E-7 : f64
%cst_118 = arith.constant 0.044194173824159216 : f64
%cst_119 = arith.constant 2.621440e+05 : f64
%cst_120 = arith.constant 0x7C00 : f16
%cst_121 = arith.constant 0x4130000000000000 : f64
%cst_122 = arith.constant 5.242880e+05 : f64
%cst_123 = arith.constant 0x4140000000000000 : f64
%c0_i64 = arith.constant 0 : i64
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%0 = tensor.empty() : tensor<1x4x128x128xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_111 : tensor<4xf16>) outs(%0 : tensor<1x4x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x4x128x128xf16>
%2 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, %cst_112 : tensor<1x4x128x128xf16>, tensor<4x4x1x1xf16>) outs(%1 : tensor<1x4x128x128xf16>) -> tensor<1x4x128x128xf16>
%padded = tensor.pad %2 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x4x128x128xf16> to tensor<?x?x?x?xf16>
%3 = tensor.empty() : tensor<1x512x128x128xf16>
%4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_109 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%5 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded, %cst_110 : tensor<?x?x?x?xf16>, tensor<512x4x3x3xf16>) outs(%4 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%collapsed = tensor.collapse_shape %5 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded = tensor.expand_shape %collapsed [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%6 = tensor.empty() : tensor<1x32x16x16384xf64>
%7 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%8 = tensor.empty() : tensor<1x32x1x1xf64>
%9 = linalg.fill ins(%cst_114 : f64) outs(%8 : tensor<1x32x1x1xf64>) -> tensor<1x32x1x1xf64>
%10 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%7 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%11 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%10 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%12 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%7, %11 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%13 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%12, %12 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%14 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%13 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%15 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%14 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%16 = tensor.empty() : tensor<1x32x1x1xf16>
%17 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%15 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%18 = linalg.fill ins(%cst_113 : f16) outs(%16 : tensor<1x32x1x1xf16>) -> tensor<1x32x1x1xf16>
%19 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%20 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%19 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%21 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%17 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%22 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%21 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%23 = tensor.empty() : tensor<1x32x16x16384xf16>
%24 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded, %20 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%25 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%24, %22 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_124 = tensor.collapse_shape %25 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_125 = tensor.expand_shape %collapsed_124 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_126 = tensor.expand_shape %cst_108 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_127 = tensor.expand_shape %cst_107 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%26 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_125, %expanded_127 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%27 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%26, %expanded_126 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%28 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%27 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%29 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%28, %27 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_128 = tensor.pad %29 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%30 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_105 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%31 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_128, %cst_106 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%30 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%collapsed_129 = tensor.collapse_shape %31 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_130 = tensor.expand_shape %collapsed_129 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%32 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_130 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%33 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%32 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%34 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%33 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%35 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%32, %34 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%36 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%35, %35 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%37 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%36 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%38 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%37 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%39 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%38 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%40 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_130 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%41 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%40 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%42 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%39 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%43 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%42 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%44 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_130, %41 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%45 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%44, %43 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_131 = tensor.collapse_shape %45 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_132 = tensor.expand_shape %collapsed_131 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_133 = tensor.expand_shape %cst_104 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_134 = tensor.expand_shape %cst_103 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%46 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_132, %expanded_134 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%47 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%46, %expanded_133 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%48 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%47 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%49 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%48, %47 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_135 = tensor.pad %49 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%50 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_101 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%51 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_135, %cst_102 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%50 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%52 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%5, %51 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%53 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%52 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%collapsed_136 = tensor.collapse_shape %53 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_137 = tensor.expand_shape %collapsed_136 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%54 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_137 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%55 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%54 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%56 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%55 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%57 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%54, %56 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%58 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%57, %57 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%59 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%58 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%60 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%59 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%61 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%60 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%62 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_137 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%63 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%62 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%64 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%61 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%65 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%64 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%66 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_137, %63 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%67 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%66, %65 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_138 = tensor.collapse_shape %67 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_139 = tensor.expand_shape %collapsed_138 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_140 = tensor.expand_shape %cst_100 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_141 = tensor.expand_shape %cst_99 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%68 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_139, %expanded_141 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%69 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%68, %expanded_140 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%collapsed_142 = tensor.collapse_shape %69 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%70 = tensor.empty() : tensor<1x16384x512xf16>
%71 = linalg.generic {indexing_maps = [#map5, #map6], iterator_types = ["parallel", "parallel", "parallel"]} ins(%collapsed_142 : tensor<1x512x16384xf16>) outs(%70 : tensor<1x16384x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x16384x512xf16>
%72 = tensor.empty() : tensor<512x512xf16>
%73 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_98 : tensor<512x512xf16>) outs(%72 : tensor<512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<512x512xf16>
%74 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%71 : tensor<1x16384x512xf16>) outs(%70 : tensor<1x16384x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x16384x512xf16>
%75 = tensor.empty() : tensor<1x512x512xf16>
%76 = linalg.generic {indexing_maps = [#map10, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%73 : tensor<512x512xf16>) outs(%75 : tensor<1x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x512xf16>
%77 = linalg.fill ins(%cst_113 : f16) outs(%70 : tensor<1x16384x512xf16>) -> tensor<1x16384x512xf16>
%78 = linalg.batch_matmul ins(%74, %76 : tensor<1x16384x512xf16>, tensor<1x512x512xf16>) outs(%77 : tensor<1x16384x512xf16>) -> tensor<1x16384x512xf16>
%79 = linalg.generic {indexing_maps = [#map9, #map11, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%78, %cst_97 : tensor<1x16384x512xf16>, tensor<512xf16>) outs(%70 : tensor<1x16384x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x16384x512xf16>
%80 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_96 : tensor<512x512xf16>) outs(%72 : tensor<512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<512x512xf16>
%81 = linalg.generic {indexing_maps = [#map10, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%80 : tensor<512x512xf16>) outs(%75 : tensor<1x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x512xf16>
%82 = linalg.batch_matmul ins(%74, %81 : tensor<1x16384x512xf16>, tensor<1x512x512xf16>) outs(%77 : tensor<1x16384x512xf16>) -> tensor<1x16384x512xf16>
%83 = linalg.generic {indexing_maps = [#map9, #map11, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%82, %cst_95 : tensor<1x16384x512xf16>, tensor<512xf16>) outs(%70 : tensor<1x16384x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x16384x512xf16>
%84 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_94 : tensor<512x512xf16>) outs(%72 : tensor<512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<512x512xf16>
%85 = linalg.generic {indexing_maps = [#map10, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%84 : tensor<512x512xf16>) outs(%75 : tensor<1x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x512xf16>
%86 = linalg.batch_matmul ins(%74, %85 : tensor<1x16384x512xf16>, tensor<1x512x512xf16>) outs(%77 : tensor<1x16384x512xf16>) -> tensor<1x16384x512xf16>
%87 = linalg.generic {indexing_maps = [#map9, #map11, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%86, %cst_93 : tensor<1x16384x512xf16>, tensor<512xf16>) outs(%70 : tensor<1x16384x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x16384x512xf16>
%expanded_143 = tensor.expand_shape %79 [[0], [1], [2, 3]] : tensor<1x16384x512xf16> into tensor<1x16384x1x512xf16>
%88 = tensor.empty() : tensor<1x1x16384x512xf16>
%89 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_143 : tensor<1x16384x1x512xf16>) outs(%88 : tensor<1x1x16384x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x1x16384x512xf16>
%collapsed_144 = tensor.collapse_shape %89 [[0], [1, 2], [3]] : tensor<1x1x16384x512xf16> into tensor<1x16384x512xf16>
%expanded_145 = tensor.expand_shape %83 [[0], [1], [2, 3]] : tensor<1x16384x512xf16> into tensor<1x16384x1x512xf16>
%90 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_145 : tensor<1x16384x1x512xf16>) outs(%88 : tensor<1x1x16384x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x1x16384x512xf16>
%collapsed_146 = tensor.collapse_shape %90 [[0], [1, 2], [3]] : tensor<1x1x16384x512xf16> into tensor<1x16384x512xf16>
%expanded_147 = tensor.expand_shape %87 [[0], [1], [2, 3]] : tensor<1x16384x512xf16> into tensor<1x16384x1x512xf16>
%91 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_147 : tensor<1x16384x1x512xf16>) outs(%88 : tensor<1x1x16384x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x1x16384x512xf16>
%collapsed_148 = tensor.collapse_shape %91 [[0], [1, 2], [3]] : tensor<1x1x16384x512xf16> into tensor<1x16384x512xf16>
%92 = tensor.empty() : tensor<1x16384x16384xf32>
%93 = tensor.empty() : tensor<1x512x16384xf16>
%94 = linalg.generic {indexing_maps = [#map5, #map6], iterator_types = ["parallel", "parallel", "parallel"]} ins(%collapsed_146 : tensor<1x16384x512xf16>) outs(%93 : tensor<1x512x16384xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x16384xf16>
%95 = tensor.empty() : tensor<1x16384x16384xf16>
%96 = linalg.fill ins(%cst_113 : f16) outs(%95 : tensor<1x16384x16384xf16>) -> tensor<1x16384x16384xf16>
%97 = linalg.batch_matmul ins(%collapsed_144, %94 : tensor<1x16384x512xf16>, tensor<1x512x16384xf16>) outs(%96 : tensor<1x16384x16384xf16>) -> tensor<1x16384x16384xf16>
%98 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%97 : tensor<1x16384x16384xf16>) outs(%95 : tensor<1x16384x16384xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_118 : f64 to f16
%589 = arith.mulf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x16384x16384xf16>
%99 = linalg.generic {indexing_maps = [#map9, #map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%98, %92 : tensor<1x16384x16384xf16>, tensor<1x16384x16384xf32>) outs(%95 : tensor<1x16384x16384xf16>) {
^bb0(%in: f16, %in_302: f32, %out: f16):
%588 = arith.truncf %in_302 : f32 to f16
%589 = arith.mulf %588, %cst_113 : f16
%590 = arith.addf %in, %589 : f16
linalg.yield %590 : f16
} -> tensor<1x16384x16384xf16>
%100 = tensor.empty() : tensor<1x16384x1xi64>
%101 = linalg.fill ins(%c0_i64 : i64) outs(%100 : tensor<1x16384x1xi64>) -> tensor<1x16384x1xi64>
%102 = tensor.empty() : tensor<1x16384x1xf16>
%103 = linalg.fill ins(%cst_116 : f16) outs(%102 : tensor<1x16384x1xf16>) -> tensor<1x16384x1xf16>
%104:2 = linalg.generic {indexing_maps = [#map5, #map13, #map13], iterator_types = ["parallel", "parallel", "reduction"]} ins(%99 : tensor<1x16384x16384xf16>) outs(%103, %101 : tensor<1x16384x1xf16>, tensor<1x16384x1xi64>) {
^bb0(%in: f16, %out: f16, %out_302: i64):
%588 = linalg.index 2 : index
%589 = arith.index_cast %588 : index to i64
%590 = arith.maxf %in, %out : f16
%591 = arith.cmpf ogt, %in, %out : f16
%592 = arith.select %591, %589, %out_302 : i64
linalg.yield %590, %592 : f16, i64
} -> (tensor<1x16384x1xf16>, tensor<1x16384x1xi64>)
%105 = linalg.generic {indexing_maps = [#map9, #map14, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%99, %104#0 : tensor<1x16384x16384xf16>, tensor<1x16384x1xf16>) outs(%95 : tensor<1x16384x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x16384x16384xf16>
%106 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%105 : tensor<1x16384x16384xf16>) outs(%95 : tensor<1x16384x16384xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.exp %in : f16
linalg.yield %588 : f16
} -> tensor<1x16384x16384xf16>
%107 = linalg.fill ins(%cst_113 : f16) outs(%102 : tensor<1x16384x1xf16>) -> tensor<1x16384x1xf16>
%108 = linalg.generic {indexing_maps = [#map5, #map13], iterator_types = ["parallel", "parallel", "reduction"]} ins(%106 : tensor<1x16384x16384xf16>) outs(%107 : tensor<1x16384x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x16384x1xf16>
%109 = linalg.generic {indexing_maps = [#map9, #map14, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%106, %108 : tensor<1x16384x16384xf16>, tensor<1x16384x1xf16>) outs(%95 : tensor<1x16384x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.divf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x16384x16384xf16>
%110 = linalg.batch_matmul ins(%109, %collapsed_148 : tensor<1x16384x16384xf16>, tensor<1x16384x512xf16>) outs(%77 : tensor<1x16384x512xf16>) -> tensor<1x16384x512xf16>
%expanded_149 = tensor.expand_shape %110 [[0], [1, 2], [3]] : tensor<1x16384x512xf16> into tensor<1x1x16384x512xf16>
%111 = tensor.empty() : tensor<1x16384x1x512xf16>
%112 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_149 : tensor<1x1x16384x512xf16>) outs(%111 : tensor<1x16384x1x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x16384x1x512xf16>
%113 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_92 : tensor<512x512xf16>) outs(%72 : tensor<512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<512x512xf16>
%collapsed_150 = tensor.collapse_shape %112 [[0, 1], [2, 3]] : tensor<1x16384x1x512xf16> into tensor<16384x512xf16>
%114 = tensor.empty() : tensor<16384x512xf16>
%115 = linalg.fill ins(%cst_113 : f16) outs(%114 : tensor<16384x512xf16>) -> tensor<16384x512xf16>
%116 = linalg.matmul ins(%collapsed_150, %113 : tensor<16384x512xf16>, tensor<512x512xf16>) outs(%115 : tensor<16384x512xf16>) -> tensor<16384x512xf16>
%117 = linalg.generic {indexing_maps = [#map15, #map7, #map7], iterator_types = ["parallel", "parallel"]} ins(%cst_91, %116 : tensor<512xf16>, tensor<16384x512xf16>) outs(%114 : tensor<16384x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<16384x512xf16>
%expanded_151 = tensor.expand_shape %117 [[0, 1], [2]] : tensor<16384x512xf16> into tensor<1x16384x512xf16>
%118 = linalg.generic {indexing_maps = [#map5, #map6], iterator_types = ["parallel", "parallel", "parallel"]} ins(%expanded_151 : tensor<1x16384x512xf16>) outs(%93 : tensor<1x512x16384xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x16384xf16>
%expanded_152 = tensor.expand_shape %118 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%119 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_152, %53 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%120 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%119 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%121 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%120 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%collapsed_153 = tensor.collapse_shape %121 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_154 = tensor.expand_shape %collapsed_153 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%122 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_154 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%123 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%122 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%124 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%123 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%125 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%122, %124 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%126 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%125, %125 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%127 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%126 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%128 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%127 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%129 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%128 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%130 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_154 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%131 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%130 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%132 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%129 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%133 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%132 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%134 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_154, %131 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%135 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%134, %133 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_155 = tensor.collapse_shape %135 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_156 = tensor.expand_shape %collapsed_155 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_157 = tensor.expand_shape %cst_90 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_158 = tensor.expand_shape %cst_89 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%136 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_156, %expanded_158 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%137 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%136, %expanded_157 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%138 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%137 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%139 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%138, %137 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_159 = tensor.pad %139 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%140 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_87 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%141 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_159, %cst_88 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%140 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%collapsed_160 = tensor.collapse_shape %141 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_161 = tensor.expand_shape %collapsed_160 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%142 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_161 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%143 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%142 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%144 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%143 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%145 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%142, %144 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%146 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%145, %145 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%147 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%146 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%148 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%147 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%149 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%148 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%150 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_161 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%151 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%150 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%152 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%149 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%153 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%152 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%154 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_161, %151 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%155 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%154, %153 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_162 = tensor.collapse_shape %155 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_163 = tensor.expand_shape %collapsed_162 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_164 = tensor.expand_shape %cst_86 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_165 = tensor.expand_shape %cst_85 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%156 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_163, %expanded_165 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%157 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%156, %expanded_164 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%158 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%157 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%159 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%158, %157 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_166 = tensor.pad %159 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%160 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_83 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%161 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_166, %cst_84 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%160 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%162 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%120, %161 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%163 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%162 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%164 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%163 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%collapsed_167 = tensor.collapse_shape %164 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_168 = tensor.expand_shape %collapsed_167 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%165 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_168 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%166 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%165 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%167 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%166 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%168 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%165, %167 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%169 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%168, %168 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%170 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%169 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%171 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%170 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%172 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%171 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%173 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_168 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%174 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%173 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%175 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%172 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%176 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%175 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%177 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_168, %174 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%178 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%177, %176 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_169 = tensor.collapse_shape %178 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_170 = tensor.expand_shape %collapsed_169 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_171 = tensor.expand_shape %cst_82 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_172 = tensor.expand_shape %cst_81 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%179 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_170, %expanded_172 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%180 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%179, %expanded_171 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%181 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%180 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%182 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%181, %180 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_173 = tensor.pad %182 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%183 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_79 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%184 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_173, %cst_80 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%183 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%collapsed_174 = tensor.collapse_shape %184 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_175 = tensor.expand_shape %collapsed_174 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%185 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_175 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%186 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%185 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%187 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%186 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%188 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%185, %187 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%189 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%188, %188 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%190 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%189 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%191 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%190 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%192 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%191 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%193 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_175 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%194 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%193 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%195 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%192 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%196 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%195 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%197 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_175, %194 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%198 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%197, %196 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_176 = tensor.collapse_shape %198 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_177 = tensor.expand_shape %collapsed_176 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_178 = tensor.expand_shape %cst_78 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_179 = tensor.expand_shape %cst_77 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%199 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_177, %expanded_179 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%200 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%199, %expanded_178 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%201 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%200 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%202 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%201, %200 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_180 = tensor.pad %202 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%203 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_75 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%204 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_180, %cst_76 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%203 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%205 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%163, %204 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%206 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%205 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%207 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%206 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%collapsed_181 = tensor.collapse_shape %207 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_182 = tensor.expand_shape %collapsed_181 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%208 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_182 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%209 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%208 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%210 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%209 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%211 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%208, %210 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%212 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%211, %211 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%213 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%212 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%214 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%213 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%215 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%214 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%216 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_182 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%217 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%216 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%218 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%215 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%219 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%218 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%220 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_182, %217 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%221 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%220, %219 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_183 = tensor.collapse_shape %221 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_184 = tensor.expand_shape %collapsed_183 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_185 = tensor.expand_shape %cst_74 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_186 = tensor.expand_shape %cst_73 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%222 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_184, %expanded_186 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%223 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%222, %expanded_185 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%224 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%223 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%225 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%224, %223 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_187 = tensor.pad %225 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%226 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_71 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%227 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_187, %cst_72 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%226 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%collapsed_188 = tensor.collapse_shape %227 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_189 = tensor.expand_shape %collapsed_188 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%228 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_189 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%229 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%228 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%230 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%229 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%231 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%228, %230 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%232 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%231, %231 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%233 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%232 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%234 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%233 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%235 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%234 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%236 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_189 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%237 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%236 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%238 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%235 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%239 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%238 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%240 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_189, %237 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%241 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%240, %239 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_190 = tensor.collapse_shape %241 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_191 = tensor.expand_shape %collapsed_190 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_192 = tensor.expand_shape %cst_70 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_193 = tensor.expand_shape %cst_69 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%242 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_191, %expanded_193 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%243 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%242, %expanded_192 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%244 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%243 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%245 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%244, %243 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_194 = tensor.pad %245 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%246 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_67 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%247 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_194, %cst_68 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%246 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%248 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%206, %247 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%249 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%248 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%250 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%249 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%collapsed_195 = tensor.collapse_shape %250 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_196 = tensor.expand_shape %collapsed_195 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%251 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_196 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%252 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%251 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%253 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%252 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%254 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%251, %253 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%255 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%254, %254 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%256 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%255 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%257 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%256 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%258 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%257 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%259 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_196 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%260 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%259 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%261 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%258 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%262 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%261 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%263 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_196, %260 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%264 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%263, %262 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_197 = tensor.collapse_shape %264 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_198 = tensor.expand_shape %collapsed_197 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_199 = tensor.expand_shape %cst_66 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_200 = tensor.expand_shape %cst_65 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%265 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_198, %expanded_200 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%266 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%265, %expanded_199 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%267 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%266 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%268 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%267, %266 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_201 = tensor.pad %268 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%269 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_63 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%270 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_201, %cst_64 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%269 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%collapsed_202 = tensor.collapse_shape %270 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
%expanded_203 = tensor.expand_shape %collapsed_202 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
%271 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_203 : tensor<1x32x16x16384xf16>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%272 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%271 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%273 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%272 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%274 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%271, %273 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%275 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%274, %274 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%6 : tensor<1x32x16x16384xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x16384xf64>
%276 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%275 : tensor<1x32x16x16384xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%277 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%276 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_119 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%278 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%277 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%279 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_203 : tensor<1x32x16x16384xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%280 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%279 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%281 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%278 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%282 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%281 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%283 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_203, %280 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%284 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%283, %282 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%23 : tensor<1x32x16x16384xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x16384xf16>
%collapsed_204 = tensor.collapse_shape %284 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
%expanded_205 = tensor.expand_shape %collapsed_204 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
%expanded_206 = tensor.expand_shape %cst_62 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_207 = tensor.expand_shape %cst_61 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%285 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_205, %expanded_207 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%286 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%285, %expanded_206 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%287 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%286 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x128x128xf16>
%288 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%287, %286 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%padded_208 = tensor.pad %288 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x128x128xf16> to tensor<?x?x?x?xf16>
%289 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_59 : tensor<512xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%290 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_208, %cst_60 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%289 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
%291 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%249, %290 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x128x128xf16>
%292 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%291 : tensor<1x512x128x128xf16>) outs(%3 : tensor<1x512x128x128xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x128x128xf16>
%293 = tensor.empty() : tensor<1x512x256x256xf16>
%294 = linalg.generic {indexing_maps = [#map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%293 : tensor<1x512x256x256xf16>) {
^bb0(%out: f16):
%588 = linalg.index 0 : index
%589 = linalg.index 1 : index
%590 = linalg.index 2 : index
%591 = linalg.index 3 : index
%592 = arith.floordivsi %590, %c2 : index
%593 = arith.floordivsi %591, %c2 : index
%extracted = tensor.extract %292[%588, %589, %592, %593] : tensor<1x512x128x128xf16>
linalg.yield %extracted : f16
} -> tensor<1x512x256x256xf16>
%padded_209 = tensor.pad %294 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x256x256xf16> to tensor<?x?x?x?xf16>
%295 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_57 : tensor<512xf16>) outs(%293 : tensor<1x512x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x512x256x256xf16>
%296 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_209, %cst_58 : tensor<?x?x?x?xf16>, tensor<512x512x3x3xf16>) outs(%295 : tensor<1x512x256x256xf16>) -> tensor<1x512x256x256xf16>
%collapsed_210 = tensor.collapse_shape %296 [[0], [1], [2, 3]] : tensor<1x512x256x256xf16> into tensor<1x512x65536xf16>
%expanded_211 = tensor.expand_shape %collapsed_210 [[0], [1, 2], [3]] : tensor<1x512x65536xf16> into tensor<1x32x16x65536xf16>
%297 = tensor.empty() : tensor<1x32x16x65536xf64>
%298 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_211 : tensor<1x32x16x65536xf16>) outs(%297 : tensor<1x32x16x65536xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x16x65536xf64>
%299 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%298 : tensor<1x32x16x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%300 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%299 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%301 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%298, %300 : tensor<1x32x16x65536xf64>, tensor<1x32x1x1xf64>) outs(%297 : tensor<1x32x16x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x65536xf64>
%302 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%301, %301 : tensor<1x32x16x65536xf64>, tensor<1x32x16x65536xf64>) outs(%297 : tensor<1x32x16x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x16x65536xf64>
%303 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%302 : tensor<1x32x16x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%304 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%303 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%305 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%304 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%306 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_211 : tensor<1x32x16x65536xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%307 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%306 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%308 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%305 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%309 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%308 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%310 = tensor.empty() : tensor<1x32x16x65536xf16>
%311 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_211, %307 : tensor<1x32x16x65536xf16>, tensor<1x32x1x1xf16>) outs(%310 : tensor<1x32x16x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x65536xf16>
%312 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%311, %309 : tensor<1x32x16x65536xf16>, tensor<1x32x1x1xf16>) outs(%310 : tensor<1x32x16x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x16x65536xf16>
%collapsed_212 = tensor.collapse_shape %312 [[0], [1, 2], [3]] : tensor<1x32x16x65536xf16> into tensor<1x512x65536xf16>
%expanded_213 = tensor.expand_shape %collapsed_212 [[0], [1], [2, 3]] : tensor<1x512x65536xf16> into tensor<1x512x256x256xf16>
%expanded_214 = tensor.expand_shape %cst_56 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%expanded_215 = tensor.expand_shape %cst_55 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
%313 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_213, %expanded_215 : tensor<1x512x256x256xf16>, tensor<1x512x1x1xf16>) outs(%293 : tensor<1x512x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x256x256xf16>
%314 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%313, %expanded_214 : tensor<1x512x256x256xf16>, tensor<1x512x1x1xf16>) outs(%293 : tensor<1x512x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x256x256xf16>
%315 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%314 : tensor<1x512x256x256xf16>) outs(%293 : tensor<1x512x256x256xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x512x256x256xf16>
%316 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%315, %314 : tensor<1x512x256x256xf16>, tensor<1x512x256x256xf16>) outs(%293 : tensor<1x512x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x512x256x256xf16>
%padded_216 = tensor.pad %316 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x512x256x256xf16> to tensor<?x?x?x?xf16>
%317 = tensor.empty() : tensor<1x256x256x256xf16>
%318 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_53 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%319 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_216, %cst_54 : tensor<?x?x?x?xf16>, tensor<256x512x3x3xf16>) outs(%318 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%collapsed_217 = tensor.collapse_shape %319 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
%expanded_218 = tensor.expand_shape %collapsed_217 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
%320 = tensor.empty() : tensor<1x32x8x65536xf64>
%321 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_218 : tensor<1x32x8x65536xf16>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%322 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%321 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%323 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%322 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%324 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%321, %323 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%325 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%324, %324 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%326 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%325 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%327 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%326 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%328 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%327 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%329 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_218 : tensor<1x32x8x65536xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%330 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%329 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%331 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%328 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%332 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%331 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%333 = tensor.empty() : tensor<1x32x8x65536xf16>
%334 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_218, %330 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%335 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%334, %332 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%collapsed_219 = tensor.collapse_shape %335 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
%expanded_220 = tensor.expand_shape %collapsed_219 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
%expanded_221 = tensor.expand_shape %cst_52 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%expanded_222 = tensor.expand_shape %cst_51 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%336 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_220, %expanded_222 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%337 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%336, %expanded_221 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%338 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%337 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x256x256x256xf16>
%339 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%338, %337 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%padded_223 = tensor.pad %339 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x256x256xf16> to tensor<?x?x?x?xf16>
%340 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_49 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%341 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_223, %cst_50 : tensor<?x?x?x?xf16>, tensor<256x256x3x3xf16>) outs(%340 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%342 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_47 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%343 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%296, %cst_48 : tensor<1x512x256x256xf16>, tensor<256x512x1x1xf16>) outs(%342 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%344 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%343, %341 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%345 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%344 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%collapsed_224 = tensor.collapse_shape %345 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
%expanded_225 = tensor.expand_shape %collapsed_224 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
%346 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_225 : tensor<1x32x8x65536xf16>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%347 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%346 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%348 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%347 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%349 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%346, %348 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%350 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%349, %349 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%351 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%350 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%352 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%351 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%353 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%352 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%354 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_225 : tensor<1x32x8x65536xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%355 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%354 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%356 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%353 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%357 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%356 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%358 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_225, %355 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%359 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%358, %357 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%collapsed_226 = tensor.collapse_shape %359 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
%expanded_227 = tensor.expand_shape %collapsed_226 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
%expanded_228 = tensor.expand_shape %cst_46 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%expanded_229 = tensor.expand_shape %cst_45 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%360 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_227, %expanded_229 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%361 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%360, %expanded_228 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%362 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%361 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x256x256x256xf16>
%363 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%362, %361 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%padded_230 = tensor.pad %363 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x256x256xf16> to tensor<?x?x?x?xf16>
%364 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_43 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%365 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_230, %cst_44 : tensor<?x?x?x?xf16>, tensor<256x256x3x3xf16>) outs(%364 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%collapsed_231 = tensor.collapse_shape %365 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
%expanded_232 = tensor.expand_shape %collapsed_231 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
%366 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_232 : tensor<1x32x8x65536xf16>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%367 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%366 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%368 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%367 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%369 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%366, %368 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%370 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%369, %369 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%371 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%370 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%372 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%371 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%373 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%372 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%374 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_232 : tensor<1x32x8x65536xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%375 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%374 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%376 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%373 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%377 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%376 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%378 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_232, %375 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%379 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%378, %377 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%collapsed_233 = tensor.collapse_shape %379 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
%expanded_234 = tensor.expand_shape %collapsed_233 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
%expanded_235 = tensor.expand_shape %cst_42 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%expanded_236 = tensor.expand_shape %cst_41 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%380 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_234, %expanded_236 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%381 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%380, %expanded_235 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%382 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%381 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x256x256x256xf16>
%383 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%382, %381 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%padded_237 = tensor.pad %383 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x256x256xf16> to tensor<?x?x?x?xf16>
%384 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_39 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%385 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_237, %cst_40 : tensor<?x?x?x?xf16>, tensor<256x256x3x3xf16>) outs(%384 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%386 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%345, %385 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%387 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%386 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%collapsed_238 = tensor.collapse_shape %387 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
%expanded_239 = tensor.expand_shape %collapsed_238 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
%388 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_239 : tensor<1x32x8x65536xf16>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%389 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%388 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%390 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%389 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%391 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%388, %390 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%392 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%391, %391 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%393 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%392 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%394 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%393 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%395 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%394 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%396 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_239 : tensor<1x32x8x65536xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%397 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%396 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%398 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%395 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%399 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%398 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%400 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_239, %397 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%401 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%400, %399 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%collapsed_240 = tensor.collapse_shape %401 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
%expanded_241 = tensor.expand_shape %collapsed_240 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
%expanded_242 = tensor.expand_shape %cst_38 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%expanded_243 = tensor.expand_shape %cst_37 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%402 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_241, %expanded_243 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%403 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%402, %expanded_242 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%404 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%403 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x256x256x256xf16>
%405 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%404, %403 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%padded_244 = tensor.pad %405 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x256x256xf16> to tensor<?x?x?x?xf16>
%406 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_35 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%407 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_244, %cst_36 : tensor<?x?x?x?xf16>, tensor<256x256x3x3xf16>) outs(%406 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%collapsed_245 = tensor.collapse_shape %407 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
%expanded_246 = tensor.expand_shape %collapsed_245 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
%408 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_246 : tensor<1x32x8x65536xf16>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%409 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%408 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%410 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%409 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%411 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%408, %410 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%412 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%411, %411 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%320 : tensor<1x32x8x65536xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x65536xf64>
%413 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%412 : tensor<1x32x8x65536xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%414 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%413 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_122 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%415 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%414 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%416 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_246 : tensor<1x32x8x65536xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%417 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%416 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%418 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%415 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%419 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%418 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%420 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_246, %417 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%421 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%420, %419 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%333 : tensor<1x32x8x65536xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x65536xf16>
%collapsed_247 = tensor.collapse_shape %421 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
%expanded_248 = tensor.expand_shape %collapsed_247 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
%expanded_249 = tensor.expand_shape %cst_34 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%expanded_250 = tensor.expand_shape %cst_33 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%422 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_248, %expanded_250 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%423 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%422, %expanded_249 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%424 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%423 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x256x256x256xf16>
%425 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%424, %423 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%padded_251 = tensor.pad %425 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x256x256xf16> to tensor<?x?x?x?xf16>
%426 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_31 : tensor<256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%427 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_251, %cst_32 : tensor<?x?x?x?xf16>, tensor<256x256x3x3xf16>) outs(%426 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
%428 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%387, %427 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x256x256xf16>
%429 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%428 : tensor<1x256x256x256xf16>) outs(%317 : tensor<1x256x256x256xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x256x256xf16>
%430 = tensor.empty() : tensor<1x256x512x512xf16>
%431 = linalg.generic {indexing_maps = [#map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%430 : tensor<1x256x512x512xf16>) {
^bb0(%out: f16):
%588 = linalg.index 0 : index
%589 = linalg.index 1 : index
%590 = linalg.index 2 : index
%591 = linalg.index 3 : index
%592 = arith.floordivsi %590, %c2 : index
%593 = arith.floordivsi %591, %c2 : index
%extracted = tensor.extract %429[%588, %589, %592, %593] : tensor<1x256x256x256xf16>
linalg.yield %extracted : f16
} -> tensor<1x256x512x512xf16>
%padded_252 = tensor.pad %431 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x512x512xf16> to tensor<?x?x?x?xf16>
%432 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_29 : tensor<256xf16>) outs(%430 : tensor<1x256x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x256x512x512xf16>
%433 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_252, %cst_30 : tensor<?x?x?x?xf16>, tensor<256x256x3x3xf16>) outs(%432 : tensor<1x256x512x512xf16>) -> tensor<1x256x512x512xf16>
%collapsed_253 = tensor.collapse_shape %433 [[0], [1], [2, 3]] : tensor<1x256x512x512xf16> into tensor<1x256x262144xf16>
%expanded_254 = tensor.expand_shape %collapsed_253 [[0], [1, 2], [3]] : tensor<1x256x262144xf16> into tensor<1x32x8x262144xf16>
%434 = tensor.empty() : tensor<1x32x8x262144xf64>
%435 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_254 : tensor<1x32x8x262144xf16>) outs(%434 : tensor<1x32x8x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x8x262144xf64>
%436 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%435 : tensor<1x32x8x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%437 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%436 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_123 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%438 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%435, %437 : tensor<1x32x8x262144xf64>, tensor<1x32x1x1xf64>) outs(%434 : tensor<1x32x8x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x262144xf64>
%439 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%438, %438 : tensor<1x32x8x262144xf64>, tensor<1x32x8x262144xf64>) outs(%434 : tensor<1x32x8x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x8x262144xf64>
%440 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%439 : tensor<1x32x8x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%441 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%440 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_123 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%442 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%441 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%443 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_254 : tensor<1x32x8x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%444 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%443 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%445 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%442 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%446 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%445 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%447 = tensor.empty() : tensor<1x32x8x262144xf16>
%448 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_254, %444 : tensor<1x32x8x262144xf16>, tensor<1x32x1x1xf16>) outs(%447 : tensor<1x32x8x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x262144xf16>
%449 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%448, %446 : tensor<1x32x8x262144xf16>, tensor<1x32x1x1xf16>) outs(%447 : tensor<1x32x8x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x8x262144xf16>
%collapsed_255 = tensor.collapse_shape %449 [[0], [1, 2], [3]] : tensor<1x32x8x262144xf16> into tensor<1x256x262144xf16>
%expanded_256 = tensor.expand_shape %collapsed_255 [[0], [1], [2, 3]] : tensor<1x256x262144xf16> into tensor<1x256x512x512xf16>
%expanded_257 = tensor.expand_shape %cst_28 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%expanded_258 = tensor.expand_shape %cst_27 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
%450 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_256, %expanded_258 : tensor<1x256x512x512xf16>, tensor<1x256x1x1xf16>) outs(%430 : tensor<1x256x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x512x512xf16>
%451 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%450, %expanded_257 : tensor<1x256x512x512xf16>, tensor<1x256x1x1xf16>) outs(%430 : tensor<1x256x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x512x512xf16>
%452 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%451 : tensor<1x256x512x512xf16>) outs(%430 : tensor<1x256x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x256x512x512xf16>
%453 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%452, %451 : tensor<1x256x512x512xf16>, tensor<1x256x512x512xf16>) outs(%430 : tensor<1x256x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x256x512x512xf16>
%padded_259 = tensor.pad %453 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x256x512x512xf16> to tensor<?x?x?x?xf16>
%454 = tensor.empty() : tensor<1x128x512x512xf16>
%455 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_25 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%456 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_259, %cst_26 : tensor<?x?x?x?xf16>, tensor<128x256x3x3xf16>) outs(%455 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%collapsed_260 = tensor.collapse_shape %456 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
%expanded_261 = tensor.expand_shape %collapsed_260 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
%457 = tensor.empty() : tensor<1x32x4x262144xf64>
%458 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_261 : tensor<1x32x4x262144xf16>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%459 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%458 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%460 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%459 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%461 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%458, %460 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%462 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%461, %461 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%463 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%462 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%464 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%463 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%465 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%464 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%466 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_261 : tensor<1x32x4x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%467 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%466 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%468 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%465 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%469 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%468 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%470 = tensor.empty() : tensor<1x32x4x262144xf16>
%471 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_261, %467 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%472 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%471, %469 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%collapsed_262 = tensor.collapse_shape %472 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
%expanded_263 = tensor.expand_shape %collapsed_262 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
%expanded_264 = tensor.expand_shape %cst_24 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%expanded_265 = tensor.expand_shape %cst_23 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%473 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_263, %expanded_265 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%474 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%473, %expanded_264 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%475 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%474 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x128x512x512xf16>
%476 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%475, %474 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%padded_266 = tensor.pad %476 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x128x512x512xf16> to tensor<?x?x?x?xf16>
%477 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_21 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%478 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_266, %cst_22 : tensor<?x?x?x?xf16>, tensor<128x128x3x3xf16>) outs(%477 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%479 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_19 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%480 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%433, %cst_20 : tensor<1x256x512x512xf16>, tensor<128x256x1x1xf16>) outs(%479 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%481 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%480, %478 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%482 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%481 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%collapsed_267 = tensor.collapse_shape %482 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
%expanded_268 = tensor.expand_shape %collapsed_267 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
%483 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_268 : tensor<1x32x4x262144xf16>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%484 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%483 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%485 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%484 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%486 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%483, %485 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%487 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%486, %486 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%488 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%487 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%489 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%488 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%490 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%489 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%491 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_268 : tensor<1x32x4x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%492 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%491 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%493 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%490 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%494 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%493 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%495 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_268, %492 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%496 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%495, %494 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%collapsed_269 = tensor.collapse_shape %496 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
%expanded_270 = tensor.expand_shape %collapsed_269 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
%expanded_271 = tensor.expand_shape %cst_18 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%expanded_272 = tensor.expand_shape %cst_17 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%497 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_270, %expanded_272 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%498 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%497, %expanded_271 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%499 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%498 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x128x512x512xf16>
%500 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%499, %498 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%padded_273 = tensor.pad %500 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x128x512x512xf16> to tensor<?x?x?x?xf16>
%501 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_15 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%502 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_273, %cst_16 : tensor<?x?x?x?xf16>, tensor<128x128x3x3xf16>) outs(%501 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%collapsed_274 = tensor.collapse_shape %502 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
%expanded_275 = tensor.expand_shape %collapsed_274 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
%503 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_275 : tensor<1x32x4x262144xf16>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%504 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%503 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%505 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%504 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%506 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%503, %505 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%507 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%506, %506 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%508 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%507 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%509 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%508 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%510 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%509 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%511 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_275 : tensor<1x32x4x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%512 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%511 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%513 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%510 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%514 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%513 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%515 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_275, %512 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%516 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%515, %514 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%collapsed_276 = tensor.collapse_shape %516 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
%expanded_277 = tensor.expand_shape %collapsed_276 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
%expanded_278 = tensor.expand_shape %cst_14 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%expanded_279 = tensor.expand_shape %cst_13 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%517 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_277, %expanded_279 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%518 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%517, %expanded_278 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%519 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%518 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x128x512x512xf16>
%520 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%519, %518 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%padded_280 = tensor.pad %520 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x128x512x512xf16> to tensor<?x?x?x?xf16>
%521 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_11 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%522 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_280, %cst_12 : tensor<?x?x?x?xf16>, tensor<128x128x3x3xf16>) outs(%521 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%523 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%482, %522 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%524 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%523 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%collapsed_281 = tensor.collapse_shape %524 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
%expanded_282 = tensor.expand_shape %collapsed_281 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
%525 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_282 : tensor<1x32x4x262144xf16>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%526 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%525 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%527 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%526 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%528 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%525, %527 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%529 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%528, %528 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%530 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%529 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%531 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%530 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%532 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%531 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%533 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_282 : tensor<1x32x4x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%534 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%533 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%535 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%532 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%536 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%535 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%537 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_282, %534 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%538 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%537, %536 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%collapsed_283 = tensor.collapse_shape %538 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
%expanded_284 = tensor.expand_shape %collapsed_283 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
%expanded_285 = tensor.expand_shape %cst_10 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%expanded_286 = tensor.expand_shape %cst_9 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%539 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_284, %expanded_286 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%540 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%539, %expanded_285 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%541 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%540 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x128x512x512xf16>
%542 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%541, %540 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%padded_287 = tensor.pad %542 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x128x512x512xf16> to tensor<?x?x?x?xf16>
%543 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_7 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%544 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_287, %cst_8 : tensor<?x?x?x?xf16>, tensor<128x128x3x3xf16>) outs(%543 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%collapsed_288 = tensor.collapse_shape %544 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
%expanded_289 = tensor.expand_shape %collapsed_288 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
%545 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_289 : tensor<1x32x4x262144xf16>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%546 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%545 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%547 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%546 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%548 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%545, %547 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%549 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%548, %548 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%550 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%549 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%551 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%550 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%552 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%551 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%553 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_289 : tensor<1x32x4x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%554 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%553 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%555 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%552 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%556 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%555 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%557 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_289, %554 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%558 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%557, %556 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%collapsed_290 = tensor.collapse_shape %558 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
%expanded_291 = tensor.expand_shape %collapsed_290 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
%expanded_292 = tensor.expand_shape %cst_6 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%expanded_293 = tensor.expand_shape %cst_5 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%559 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_291, %expanded_293 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%560 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%559, %expanded_292 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%561 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%560 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x128x512x512xf16>
%562 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%561, %560 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%padded_294 = tensor.pad %562 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x128x512x512xf16> to tensor<?x?x?x?xf16>
%563 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_3 : tensor<128xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%564 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_294, %cst_4 : tensor<?x?x?x?xf16>, tensor<128x128x3x3xf16>) outs(%563 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
%565 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%524, %564 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%566 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%565 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x128x512x512xf16>
%collapsed_295 = tensor.collapse_shape %566 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
%expanded_296 = tensor.expand_shape %collapsed_295 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
%567 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_296 : tensor<1x32x4x262144xf16>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f16, %out: f64):
%588 = arith.extf %in : f16 to f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%568 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%567 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%569 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%568 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%570 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%567, %569 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.subf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%571 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%570, %570 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%457 : tensor<1x32x4x262144xf64>) {
^bb0(%in: f64, %in_302: f64, %out: f64):
%588 = arith.mulf %in, %in_302 : f64
linalg.yield %588 : f64
} -> tensor<1x32x4x262144xf64>
%572 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%571 : tensor<1x32x4x262144xf64>) outs(%9 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.addf %in, %out : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%573 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%572 : tensor<1x32x1x1xf64>) outs(%8 : tensor<1x32x1x1xf64>) {
^bb0(%in: f64, %out: f64):
%588 = arith.divf %in, %cst_121 : f64
linalg.yield %588 : f64
} -> tensor<1x32x1x1xf64>
%574 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%573 : tensor<1x32x1x1xf64>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f64, %out: f16):
%588 = arith.truncf %in : f64 to f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%575 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%expanded_296 : tensor<1x32x4x262144xf16>) outs(%18 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.addf %in, %out : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%576 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%575 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.divf %in, %cst_120 : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%577 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%574 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.truncf %cst_117 : f64 to f16
%589 = arith.addf %in, %588 : f16
linalg.yield %589 : f16
} -> tensor<1x32x1x1xf16>
%578 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%577 : tensor<1x32x1x1xf16>) outs(%16 : tensor<1x32x1x1xf16>) {
^bb0(%in: f16, %out: f16):
%588 = math.rsqrt %in : f16
linalg.yield %588 : f16
} -> tensor<1x32x1x1xf16>
%579 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_296, %576 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.subf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%580 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%579, %578 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%470 : tensor<1x32x4x262144xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x32x4x262144xf16>
%collapsed_297 = tensor.collapse_shape %580 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
%expanded_298 = tensor.expand_shape %collapsed_297 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
%expanded_299 = tensor.expand_shape %cst_2 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%expanded_300 = tensor.expand_shape %cst_1 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
%581 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_298, %expanded_300 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%582 = linalg.generic {indexing_maps = [#map2, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%581, %expanded_299 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.addf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%583 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%582 : tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %out: f16):
%588 = arith.negf %in : f16
%589 = math.exp %588 : f16
%590 = arith.addf %589, %cst_115 : f16
%591 = arith.divf %cst_115, %590 : f16
linalg.yield %591 : f16
} -> tensor<1x128x512x512xf16>
%584 = linalg.generic {indexing_maps = [#map2, #map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%583, %582 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%454 : tensor<1x128x512x512xf16>) {
^bb0(%in: f16, %in_302: f16, %out: f16):
%588 = arith.mulf %in, %in_302 : f16
linalg.yield %588 : f16
} -> tensor<1x128x512x512xf16>
%padded_301 = tensor.pad %584 low[%c0, %c0, %c1, %c1] high[%c0, %c0, %c1, %c1] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_113 : f16
} : tensor<1x128x512x512xf16> to tensor<?x?x?x?xf16>
%585 = tensor.empty() : tensor<1x3x512x512xf16>
%586 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst : tensor<3xf16>) outs(%585 : tensor<1x3x512x512xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<1x3x512x512xf16>
%587 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_301, %cst_0 : tensor<?x?x?x?xf16>, tensor<3x128x3x3xf16>) outs(%586 : tensor<1x3x512x512xf16>) -> tensor<1x3x512x512xf16>
return %587 : tensor<1x3x512x512xf16>
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment