Example
transformer_ad.kina
Source
let fwd = fn(
x: tensor<F32, [1, 32, 64], @host>,
gamma1: tensor<F32, [1, 1, 64], @host>,
wq: tensor<F32, [1, 64, 64], @host>,
wk: tensor<F32, [1, 64, 64], @host>,
wv: tensor<F32, [1, 64, 64], @host>,
wo: tensor<F32, [1, 64, 64], @host>,
gamma2: tensor<F32, [1, 1, 64], @host>,
w1: tensor<F32, [1, 64, 128], @host>,
w2: tensor<F32, [1, 128, 64], @host>
) -> tensor<F32, [1, 32, 64], @host> {
// rmsnorm 1
let x_sq = x * x in
let sum_sq = x_sq.reduce(op=sum, axis=2, keepdims=true) in
let d_inv = tensor<F32, [1, 1, 1], @host> { 0.015625 } in
let d_inv_bcast1 = d_inv.broadcast([1, 32, 1]) in
let mean_sq = sum_sq * d_inv_bcast1 in
let eps = tensor<F32, [1, 1, 1], @host> { 0.00001 } in
let eps_bcast = eps.broadcast([1, 32, 1]) in
let variance = mean_sq + eps_bcast in
let inv_std = rsqrt(variance) in
let inv_std_bcast = inv_std.broadcast([1, 32, 64]) in
let x_norm1 = x * inv_std_bcast in
let gamma_bcast1 = gamma1.broadcast([1, 32, 64]) in
let norm1 = x_norm1 * gamma_bcast1 in
// attention
let q = norm1 @ wq in
let k = norm1 @ wk in
let v = norm1 @ wv in
let k_t = k.transpose([0, 2, 1]) in
let scores = q @ k_t in
let scale = tensor<F32, [1, 1, 1], @host> { 0.125 } in
let scale_bcast = scale.broadcast([1, 32, 32]) in
let scores_scaled = scores * scale_bcast in
let probs = scores_scaled.softmax(axis=2) in
let attn_out = probs @ v in
let attn = attn_out @ wo in
// res 1
let res1 = x + attn in
// rmsnorm 2
let res1_sq = res1 * res1 in
let res1_sum_sq = res1_sq.reduce(op=sum, axis=2, keepdims=true) in
let res1_mean_sq = res1_sum_sq * d_inv_bcast1 in
let res1_variance = res1_mean_sq + eps_bcast in
let res1_inv_std = rsqrt(res1_variance) in
let res1_inv_std_bcast = res1_inv_std.broadcast([1, 32, 64]) in
let x_norm2 = res1 * res1_inv_std_bcast in
let gamma_bcast2 = gamma2.broadcast([1, 32, 64]) in
let norm2 = x_norm2 * gamma_bcast2 in
// swiglu
let h_ffn = norm2 @ w1 in
let h_sig = sigmoid(h_ffn) in
let h_swish = h_ffn * h_sig in
let ffn = h_swish @ w2 in
// res 2
let res2 = res1 + ffn in
res2
} in
let x_input = tensor<F32, [1, 32, 64], @host> { 1.0 } in
let gamma1 = tensor<F32, [1, 1, 64], @host> { 1.0 } in
let wq = tensor<F32, [1, 64, 64], @host> { 0.1 } in
let wk = tensor<F32, [1, 64, 64], @host> { 0.1 } in
let wv = tensor<F32, [1, 64, 64], @host> { 0.1 } in
let wo = tensor<F32, [1, 64, 64], @host> { 0.1 } in
let gamma2 = tensor<F32, [1, 1, 64], @host> { 1.0 } in
let w1 = tensor<F32, [1, 64, 128], @host> { 0.1 } in
let w2 = tensor<F32, [1, 128, 64], @host> { 0.1 } in
handle fwd(x_input, gamma1, wq, wk, wv, wo, gamma2, w1, w2) with reverseMode
Emitted MLIR (golden)
module {
func.func @fn_val_0(%1: tensor<1x32x64xf32>, %2: tensor<1x1x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%3 = tensor.empty() : tensor<1x32x64xf32>
%4 = linalg.mul ins(%1, %1 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%3 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%7 = arith.constant 0.000000e+00 : f32
%5 = tensor.empty() : tensor<1x32x1xf32>
%6 = linalg.fill ins(%7 : f32) outs(%5 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%8 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%4 : tensor<1x32x64xf32>) outs(%6 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%9 = arith.addf %a0, %a1 : f32
linalg.yield %9 : f32
} -> tensor<1x32x1xf32>
%10 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%11 = tensor.empty() : tensor<1x32x1xf32>
%12 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%10 : tensor<1x1x1xf32>) outs(%11 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%13 = tensor.empty() : tensor<1x32x1xf32>
%14 = linalg.mul ins(%8, %12 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%13 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%15 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%16 = tensor.empty() : tensor<1x32x1xf32>
%17 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%15 : tensor<1x1x1xf32>) outs(%16 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%18 = tensor.empty() : tensor<1x32x1xf32>
%19 = linalg.add ins(%14, %17 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%18 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%23 = tensor.empty() : tensor<1x32x1xf32>
%24 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%19 : tensor<1x32x1xf32>) outs(%23 : tensor<1x32x1xf32>) {
^bb0(%x: f32, %_: f32):
%20 = math.sqrt %x : f32
%21 = arith.constant 1.000000e+00 : f32
%22 = arith.divf %21, %20 : f32
linalg.yield %22 : f32
} -> tensor<1x32x1xf32>
%25 = tensor.empty() : tensor<1x32x64xf32>
%26 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%24 : tensor<1x32x1xf32>) outs(%25 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%27 = tensor.empty() : tensor<1x32x64xf32>
%28 = linalg.mul ins(%1, %26 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%27 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%29 = tensor.empty() : tensor<1x32x64xf32>
%30 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%2 : tensor<1x1x64xf32>) outs(%29 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%31 = tensor.empty() : tensor<1x32x64xf32>
%32 = linalg.mul ins(%28, %30 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%31 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %32 : tensor<1x32x64xf32>
}
func.func @fn_val_33(%34: tensor<1x32x64xf32>, %35: tensor<1x64x128xf32>, %36: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%37 = tensor.empty() : tensor<1x32x128xf32>
%38 = linalg.batch_matmul ins(%34, %35 : tensor<1x32x64xf32>, tensor<1x64x128xf32>) outs(%37 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%44 = tensor.empty() : tensor<1x32x128xf32>
%45 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%38 : tensor<1x32x128xf32>) outs(%44 : tensor<1x32x128xf32>) {
^bb0(%a0: f32, %a1: f32):
%39 = arith.constant 1.000000e+00 : f32
%40 = arith.negf %a0 : f32
%41 = math.exp %40 : f32
%42 = arith.addf %39, %41 : f32
%43 = arith.divf %39, %42 : f32
linalg.yield %43 : f32
} -> tensor<1x32x128xf32>
%46 = tensor.empty() : tensor<1x32x128xf32>
%47 = linalg.mul ins(%38, %45 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%46 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%48 = tensor.empty() : tensor<1x32x64xf32>
%49 = linalg.batch_matmul ins(%47, %36 : tensor<1x32x128xf32>, tensor<1x128x64xf32>) outs(%48 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %49 : tensor<1x32x64xf32>
}
func.func @fn_val_50(%51: tensor<1x32x64xf32>, %52: tensor<1x64x64xf32>, %53: tensor<1x64x64xf32>, %54: tensor<1x64x64xf32>, %55: tensor<1x64x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%56 = tensor.empty() : tensor<1x32x64xf32>
%57 = linalg.batch_matmul ins(%51, %52 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%56 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%58 = tensor.empty() : tensor<1x32x64xf32>
%59 = linalg.batch_matmul ins(%51, %53 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%58 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%60 = tensor.empty() : tensor<1x32x64xf32>
%61 = linalg.batch_matmul ins(%51, %54 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%60 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%62 = tensor.empty() : tensor<1x64x32xf32>
%63 = linalg.transpose ins(%59 : tensor<1x32x64xf32>) outs(%62 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%64 = tensor.empty() : tensor<1x32x32xf32>
%65 = linalg.batch_matmul ins(%57, %63 : tensor<1x32x64xf32>, tensor<1x64x32xf32>) outs(%64 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%66 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%67 = tensor.empty() : tensor<1x32x32xf32>
%68 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%66 : tensor<1x1x1xf32>) outs(%67 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%69 = tensor.empty() : tensor<1x32x32xf32>
%70 = linalg.mul ins(%65, %68 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%69 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%73 = arith.constant 0xFF800000 : f32
%71 = tensor.empty() : tensor<1x32x1xf32>
%72 = linalg.fill ins(%73 : f32) outs(%71 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%74 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%70 : tensor<1x32x32xf32>) outs(%72 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%75 = arith.maximumf %a0, %a1 : f32
linalg.yield %75 : f32
} -> tensor<1x32x1xf32>
%76 = tensor.empty() : tensor<1x32x32xf32>
%77 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%74 : tensor<1x32x1xf32>) outs(%76 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%78 = tensor.empty() : tensor<1x32x32xf32>
%79 = linalg.sub ins(%70, %77 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%78 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%80 = tensor.empty() : tensor<1x32x32xf32>
%81 = linalg.exp ins(%79 : tensor<1x32x32xf32>) outs(%80 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%84 = arith.constant 0xFF800000 : f32
%82 = tensor.empty() : tensor<1x32x1xf32>
%83 = linalg.fill ins(%84 : f32) outs(%82 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%85 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%70 : tensor<1x32x32xf32>) outs(%83 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%86 = arith.maximumf %a0, %a1 : f32
linalg.yield %86 : f32
} -> tensor<1x32x1xf32>
%87 = tensor.empty() : tensor<1x32x32xf32>
%88 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%85 : tensor<1x32x1xf32>) outs(%87 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%89 = tensor.empty() : tensor<1x32x32xf32>
%90 = linalg.sub ins(%70, %88 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%89 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%91 = tensor.empty() : tensor<1x32x32xf32>
%92 = linalg.exp ins(%90 : tensor<1x32x32xf32>) outs(%91 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%95 = arith.constant 0.000000e+00 : f32
%93 = tensor.empty() : tensor<1x32x1xf32>
%94 = linalg.fill ins(%95 : f32) outs(%93 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%96 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%92 : tensor<1x32x32xf32>) outs(%94 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%97 = arith.addf %a0, %a1 : f32
linalg.yield %97 : f32
} -> tensor<1x32x1xf32>
%98 = tensor.empty() : tensor<1x32x32xf32>
%99 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%96 : tensor<1x32x1xf32>) outs(%98 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%100 = tensor.empty() : tensor<1x32x32xf32>
%101 = linalg.div ins(%81, %99 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%100 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%102 = tensor.empty() : tensor<1x32x64xf32>
%103 = linalg.batch_matmul ins(%101, %61 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%102 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%104 = tensor.empty() : tensor<1x32x64xf32>
%105 = linalg.batch_matmul ins(%103, %55 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%104 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %105 : tensor<1x32x64xf32>
}
func.func @fn_val_106(%107: tensor<1x32x64xf32>, %108: tensor<1x1x64xf32>, %109: tensor<1x64x64xf32>, %110: tensor<1x64x64xf32>, %111: tensor<1x64x64xf32>, %112: tensor<1x64x64xf32>, %113: tensor<1x1x64xf32>, %114: tensor<1x64x128xf32>, %115: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%116 = func.call @fn_val_0(%107, %108) : (tensor<1x32x64xf32>, tensor<1x1x64xf32>) -> tensor<1x32x64xf32>
%117 = func.call @fn_val_50(%116, %109, %110, %111, %112) : (tensor<1x32x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>) -> tensor<1x32x64xf32>
%118 = tensor.empty() : tensor<1x32x64xf32>
%119 = linalg.add ins(%107, %117 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%118 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%120 = func.call @fn_val_0(%119, %113) : (tensor<1x32x64xf32>, tensor<1x1x64xf32>) -> tensor<1x32x64xf32>
%121 = func.call @fn_val_33(%120, %114, %115) : (tensor<1x32x64xf32>, tensor<1x64x128xf32>, tensor<1x128x64xf32>) -> tensor<1x32x64xf32>
%122 = tensor.empty() : tensor<1x32x64xf32>
%123 = linalg.add ins(%119, %121 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%122 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %123 : tensor<1x32x64xf32>
}
func.func @fn_val_124(%125: tensor<1x32x64xf32>, %126: tensor<1x1x64xf32>, %127: tensor<1x64x64xf32>, %128: tensor<1x64x64xf32>, %129: tensor<1x64x64xf32>, %130: tensor<1x64x64xf32>, %131: tensor<1x1x64xf32>, %132: tensor<1x64x128xf32>, %133: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%134 = tensor.empty() : tensor<1x32x64xf32>
%135 = linalg.mul ins(%125, %125 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%134 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%138 = arith.constant 0.000000e+00 : f32
%136 = tensor.empty() : tensor<1x32x1xf32>
%137 = linalg.fill ins(%138 : f32) outs(%136 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%139 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%135 : tensor<1x32x64xf32>) outs(%137 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%140 = arith.addf %a0, %a1 : f32
linalg.yield %140 : f32
} -> tensor<1x32x1xf32>
%141 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%142 = tensor.empty() : tensor<1x32x1xf32>
%143 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%141 : tensor<1x1x1xf32>) outs(%142 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%144 = tensor.empty() : tensor<1x32x1xf32>
%145 = linalg.mul ins(%139, %143 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%144 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%146 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%147 = tensor.empty() : tensor<1x32x1xf32>
%148 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%146 : tensor<1x1x1xf32>) outs(%147 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%149 = tensor.empty() : tensor<1x32x1xf32>
%150 = linalg.add ins(%145, %148 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%149 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%154 = tensor.empty() : tensor<1x32x1xf32>
%155 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%150 : tensor<1x32x1xf32>) outs(%154 : tensor<1x32x1xf32>) {
^bb0(%x: f32, %_: f32):
%151 = math.sqrt %x : f32
%152 = arith.constant 1.000000e+00 : f32
%153 = arith.divf %152, %151 : f32
linalg.yield %153 : f32
} -> tensor<1x32x1xf32>
%156 = tensor.empty() : tensor<1x32x64xf32>
%157 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%155 : tensor<1x32x1xf32>) outs(%156 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%158 = tensor.empty() : tensor<1x32x64xf32>
%159 = linalg.mul ins(%125, %157 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%158 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%160 = tensor.empty() : tensor<1x32x64xf32>
%161 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%126 : tensor<1x1x64xf32>) outs(%160 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%162 = tensor.empty() : tensor<1x32x64xf32>
%163 = linalg.mul ins(%159, %161 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%162 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%164 = tensor.empty() : tensor<1x32x64xf32>
%165 = linalg.batch_matmul ins(%163, %127 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%164 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%166 = tensor.empty() : tensor<1x32x64xf32>
%167 = linalg.batch_matmul ins(%163, %128 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%166 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%168 = tensor.empty() : tensor<1x32x64xf32>
%169 = linalg.batch_matmul ins(%163, %129 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%168 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%170 = tensor.empty() : tensor<1x64x32xf32>
%171 = linalg.transpose ins(%167 : tensor<1x32x64xf32>) outs(%170 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%172 = tensor.empty() : tensor<1x32x32xf32>
%173 = linalg.batch_matmul ins(%165, %171 : tensor<1x32x64xf32>, tensor<1x64x32xf32>) outs(%172 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%174 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%175 = tensor.empty() : tensor<1x32x32xf32>
%176 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%174 : tensor<1x1x1xf32>) outs(%175 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%177 = tensor.empty() : tensor<1x32x32xf32>
%178 = linalg.mul ins(%173, %176 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%177 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%181 = arith.constant 0xFF800000 : f32
%179 = tensor.empty() : tensor<1x32x1xf32>
%180 = linalg.fill ins(%181 : f32) outs(%179 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%182 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%178 : tensor<1x32x32xf32>) outs(%180 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%183 = arith.maximumf %a0, %a1 : f32
linalg.yield %183 : f32
} -> tensor<1x32x1xf32>
%184 = tensor.empty() : tensor<1x32x32xf32>
%185 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%182 : tensor<1x32x1xf32>) outs(%184 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%186 = tensor.empty() : tensor<1x32x32xf32>
%187 = linalg.sub ins(%178, %185 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%186 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%188 = tensor.empty() : tensor<1x32x32xf32>
%189 = linalg.exp ins(%187 : tensor<1x32x32xf32>) outs(%188 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%192 = arith.constant 0xFF800000 : f32
%190 = tensor.empty() : tensor<1x32x1xf32>
%191 = linalg.fill ins(%192 : f32) outs(%190 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%193 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%178 : tensor<1x32x32xf32>) outs(%191 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%194 = arith.maximumf %a0, %a1 : f32
linalg.yield %194 : f32
} -> tensor<1x32x1xf32>
%195 = tensor.empty() : tensor<1x32x32xf32>
%196 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%193 : tensor<1x32x1xf32>) outs(%195 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%197 = tensor.empty() : tensor<1x32x32xf32>
%198 = linalg.sub ins(%178, %196 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%197 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%199 = tensor.empty() : tensor<1x32x32xf32>
%200 = linalg.exp ins(%198 : tensor<1x32x32xf32>) outs(%199 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%203 = arith.constant 0.000000e+00 : f32
%201 = tensor.empty() : tensor<1x32x1xf32>
%202 = linalg.fill ins(%203 : f32) outs(%201 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%204 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%200 : tensor<1x32x32xf32>) outs(%202 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%205 = arith.addf %a0, %a1 : f32
linalg.yield %205 : f32
} -> tensor<1x32x1xf32>
%206 = tensor.empty() : tensor<1x32x32xf32>
%207 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%204 : tensor<1x32x1xf32>) outs(%206 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%208 = tensor.empty() : tensor<1x32x32xf32>
%209 = linalg.div ins(%189, %207 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%208 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%210 = tensor.empty() : tensor<1x32x64xf32>
%211 = linalg.batch_matmul ins(%209, %169 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%210 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%212 = tensor.empty() : tensor<1x32x64xf32>
%213 = linalg.batch_matmul ins(%211, %130 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%212 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%214 = tensor.empty() : tensor<1x32x64xf32>
%215 = linalg.add ins(%125, %213 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%214 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%216 = tensor.empty() : tensor<1x32x64xf32>
%217 = linalg.mul ins(%215, %215 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%216 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%220 = arith.constant 0.000000e+00 : f32
%218 = tensor.empty() : tensor<1x32x1xf32>
%219 = linalg.fill ins(%220 : f32) outs(%218 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%221 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%217 : tensor<1x32x64xf32>) outs(%219 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%222 = arith.addf %a0, %a1 : f32
linalg.yield %222 : f32
} -> tensor<1x32x1xf32>
%223 = tensor.empty() : tensor<1x32x1xf32>
%224 = linalg.mul ins(%221, %143 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%223 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%225 = tensor.empty() : tensor<1x32x1xf32>
%226 = linalg.add ins(%224, %148 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%225 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%230 = tensor.empty() : tensor<1x32x1xf32>
%231 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%226 : tensor<1x32x1xf32>) outs(%230 : tensor<1x32x1xf32>) {
^bb0(%x: f32, %_: f32):
%227 = math.sqrt %x : f32
%228 = arith.constant 1.000000e+00 : f32
%229 = arith.divf %228, %227 : f32
linalg.yield %229 : f32
} -> tensor<1x32x1xf32>
%232 = tensor.empty() : tensor<1x32x64xf32>
%233 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%231 : tensor<1x32x1xf32>) outs(%232 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%234 = tensor.empty() : tensor<1x32x64xf32>
%235 = linalg.mul ins(%215, %233 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%234 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%236 = tensor.empty() : tensor<1x32x64xf32>
%237 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%131 : tensor<1x1x64xf32>) outs(%236 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%238 = tensor.empty() : tensor<1x32x64xf32>
%239 = linalg.mul ins(%235, %237 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%238 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%240 = tensor.empty() : tensor<1x32x128xf32>
%241 = linalg.batch_matmul ins(%239, %132 : tensor<1x32x64xf32>, tensor<1x64x128xf32>) outs(%240 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%247 = tensor.empty() : tensor<1x32x128xf32>
%248 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%241 : tensor<1x32x128xf32>) outs(%247 : tensor<1x32x128xf32>) {
^bb0(%a0: f32, %a1: f32):
%242 = arith.constant 1.000000e+00 : f32
%243 = arith.negf %a0 : f32
%244 = math.exp %243 : f32
%245 = arith.addf %242, %244 : f32
%246 = arith.divf %242, %245 : f32
linalg.yield %246 : f32
} -> tensor<1x32x128xf32>
%249 = tensor.empty() : tensor<1x32x128xf32>
%250 = linalg.mul ins(%241, %248 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%249 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%251 = tensor.empty() : tensor<1x32x64xf32>
%252 = linalg.batch_matmul ins(%250, %133 : tensor<1x32x128xf32>, tensor<1x128x64xf32>) outs(%251 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%253 = tensor.empty() : tensor<1x32x64xf32>
%254 = linalg.add ins(%215, %252 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%253 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %254 : tensor<1x32x64xf32>
}
func.func @main() {
%255 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%256 = arith.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%257 = arith.constant dense<0.000000e+00> : tensor<1x64x64xf32>
%258 = arith.constant dense<0.000000e+00> : tensor<1x64x64xf32>
%259 = arith.constant dense<0.000000e+00> : tensor<1x64x64xf32>
%260 = arith.constant dense<0.000000e+00> : tensor<1x64x64xf32>
%261 = arith.constant dense<0.000000e+00> : tensor<1x1x64xf32>
%262 = arith.constant dense<0.000000e+00> : tensor<1x64x128xf32>
%263 = arith.constant dense<0.000000e+00> : tensor<1x128x64xf32>
%264 = tensor.empty() : tensor<1x32x64xf32>
%265 = linalg.mul ins(%255, %255 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%264 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%268 = arith.constant 0.000000e+00 : f32
%266 = tensor.empty() : tensor<1x32x1xf32>
%267 = linalg.fill ins(%268 : f32) outs(%266 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%269 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%265 : tensor<1x32x64xf32>) outs(%267 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%270 = arith.addf %a0, %a1 : f32
linalg.yield %270 : f32
} -> tensor<1x32x1xf32>
%271 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%272 = tensor.empty() : tensor<1x32x1xf32>
%273 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%271 : tensor<1x1x1xf32>) outs(%272 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%274 = tensor.empty() : tensor<1x32x1xf32>
%275 = linalg.mul ins(%269, %273 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%274 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%276 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%277 = tensor.empty() : tensor<1x32x1xf32>
%278 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%276 : tensor<1x1x1xf32>) outs(%277 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%279 = tensor.empty() : tensor<1x32x1xf32>
%280 = linalg.add ins(%275, %278 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%279 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%284 = tensor.empty() : tensor<1x32x1xf32>
%285 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%280 : tensor<1x32x1xf32>) outs(%284 : tensor<1x32x1xf32>) {
^bb0(%x: f32, %_: f32):
%281 = math.sqrt %x : f32
%282 = arith.constant 1.000000e+00 : f32
%283 = arith.divf %282, %281 : f32
linalg.yield %283 : f32
} -> tensor<1x32x1xf32>
%286 = tensor.empty() : tensor<1x32x64xf32>
%287 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%285 : tensor<1x32x1xf32>) outs(%286 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%288 = tensor.empty() : tensor<1x32x64xf32>
%289 = linalg.mul ins(%255, %287 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%288 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%290 = tensor.empty() : tensor<1x32x64xf32>
%291 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%256 : tensor<1x1x64xf32>) outs(%290 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%292 = tensor.empty() : tensor<1x32x64xf32>
%293 = linalg.mul ins(%289, %291 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%292 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%294 = tensor.empty() : tensor<1x32x64xf32>
%295 = linalg.batch_matmul ins(%293, %257 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%294 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%296 = tensor.empty() : tensor<1x32x64xf32>
%297 = linalg.batch_matmul ins(%293, %258 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%296 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%298 = tensor.empty() : tensor<1x32x64xf32>
%299 = linalg.batch_matmul ins(%293, %259 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%298 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%300 = tensor.empty() : tensor<1x64x32xf32>
%301 = linalg.transpose ins(%297 : tensor<1x32x64xf32>) outs(%300 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%302 = tensor.empty() : tensor<1x32x32xf32>
%303 = linalg.batch_matmul ins(%295, %301 : tensor<1x32x64xf32>, tensor<1x64x32xf32>) outs(%302 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%304 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%305 = tensor.empty() : tensor<1x32x32xf32>
%306 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%304 : tensor<1x1x1xf32>) outs(%305 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%307 = tensor.empty() : tensor<1x32x32xf32>
%308 = linalg.mul ins(%303, %306 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%307 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%311 = arith.constant 0xFF800000 : f32
%309 = tensor.empty() : tensor<1x32x1xf32>
%310 = linalg.fill ins(%311 : f32) outs(%309 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%312 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%308 : tensor<1x32x32xf32>) outs(%310 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%313 = arith.maximumf %a0, %a1 : f32
linalg.yield %313 : f32
} -> tensor<1x32x1xf32>
%314 = tensor.empty() : tensor<1x32x32xf32>
%315 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%312 : tensor<1x32x1xf32>) outs(%314 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%316 = tensor.empty() : tensor<1x32x32xf32>
%317 = linalg.sub ins(%308, %315 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%316 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%318 = tensor.empty() : tensor<1x32x32xf32>
%319 = linalg.exp ins(%317 : tensor<1x32x32xf32>) outs(%318 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%322 = arith.constant 0xFF800000 : f32
%320 = tensor.empty() : tensor<1x32x1xf32>
%321 = linalg.fill ins(%322 : f32) outs(%320 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%323 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%308 : tensor<1x32x32xf32>) outs(%321 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%324 = arith.maximumf %a0, %a1 : f32
linalg.yield %324 : f32
} -> tensor<1x32x1xf32>
%325 = tensor.empty() : tensor<1x32x32xf32>
%326 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%323 : tensor<1x32x1xf32>) outs(%325 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%327 = tensor.empty() : tensor<1x32x32xf32>
%328 = linalg.sub ins(%308, %326 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%327 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%329 = tensor.empty() : tensor<1x32x32xf32>
%330 = linalg.exp ins(%328 : tensor<1x32x32xf32>) outs(%329 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%333 = arith.constant 0.000000e+00 : f32
%331 = tensor.empty() : tensor<1x32x1xf32>
%332 = linalg.fill ins(%333 : f32) outs(%331 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%334 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%330 : tensor<1x32x32xf32>) outs(%332 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%335 = arith.addf %a0, %a1 : f32
linalg.yield %335 : f32
} -> tensor<1x32x1xf32>
%336 = tensor.empty() : tensor<1x32x32xf32>
%337 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%334 : tensor<1x32x1xf32>) outs(%336 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%338 = tensor.empty() : tensor<1x32x32xf32>
%339 = linalg.div ins(%319, %337 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%338 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%340 = tensor.empty() : tensor<1x32x64xf32>
%341 = linalg.batch_matmul ins(%339, %299 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%340 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%342 = tensor.empty() : tensor<1x32x64xf32>
%343 = linalg.batch_matmul ins(%341, %260 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%342 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%344 = tensor.empty() : tensor<1x32x64xf32>
%345 = linalg.add ins(%255, %343 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%344 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%346 = tensor.empty() : tensor<1x32x64xf32>
%347 = linalg.mul ins(%345, %345 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%346 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%350 = arith.constant 0.000000e+00 : f32
%348 = tensor.empty() : tensor<1x32x1xf32>
%349 = linalg.fill ins(%350 : f32) outs(%348 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%351 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%347 : tensor<1x32x64xf32>) outs(%349 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%352 = arith.addf %a0, %a1 : f32
linalg.yield %352 : f32
} -> tensor<1x32x1xf32>
%353 = tensor.empty() : tensor<1x32x1xf32>
%354 = linalg.mul ins(%351, %273 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%353 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%355 = tensor.empty() : tensor<1x32x1xf32>
%356 = linalg.add ins(%354, %278 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%355 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%360 = tensor.empty() : tensor<1x32x1xf32>
%361 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%356 : tensor<1x32x1xf32>) outs(%360 : tensor<1x32x1xf32>) {
^bb0(%x: f32, %_: f32):
%357 = math.sqrt %x : f32
%358 = arith.constant 1.000000e+00 : f32
%359 = arith.divf %358, %357 : f32
linalg.yield %359 : f32
} -> tensor<1x32x1xf32>
%362 = tensor.empty() : tensor<1x32x64xf32>
%363 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%361 : tensor<1x32x1xf32>) outs(%362 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%364 = tensor.empty() : tensor<1x32x64xf32>
%365 = linalg.mul ins(%345, %363 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%364 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%366 = tensor.empty() : tensor<1x32x64xf32>
%367 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%261 : tensor<1x1x64xf32>) outs(%366 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%368 = tensor.empty() : tensor<1x32x64xf32>
%369 = linalg.mul ins(%365, %367 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%368 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%370 = tensor.empty() : tensor<1x32x128xf32>
%371 = linalg.batch_matmul ins(%369, %262 : tensor<1x32x64xf32>, tensor<1x64x128xf32>) outs(%370 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%377 = tensor.empty() : tensor<1x32x128xf32>
%378 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%371 : tensor<1x32x128xf32>) outs(%377 : tensor<1x32x128xf32>) {
^bb0(%a0: f32, %a1: f32):
%372 = arith.constant 1.000000e+00 : f32
%373 = arith.negf %a0 : f32
%374 = math.exp %373 : f32
%375 = arith.addf %372, %374 : f32
%376 = arith.divf %372, %375 : f32
linalg.yield %376 : f32
} -> tensor<1x32x128xf32>
%379 = tensor.empty() : tensor<1x32x128xf32>
%380 = linalg.mul ins(%371, %378 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%379 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%381 = tensor.empty() : tensor<1x32x64xf32>
%382 = linalg.batch_matmul ins(%380, %263 : tensor<1x32x128xf32>, tensor<1x128x64xf32>) outs(%381 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%383 = tensor.empty() : tensor<1x32x64xf32>
%384 = linalg.add ins(%345, %382 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%383 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%385 = arith.constant dense<1.000000e+00> : tensor<1x32x64xf32>
%386 = tensor.empty() : tensor<1x128x32xf32>
%387 = linalg.transpose ins(%380 : tensor<1x32x128xf32>) outs(%386 : tensor<1x128x32xf32>) permutation = [0, 2, 1]
%388 = arith.constant dense<1.000000e+00> : tensor<1x32x64xf32>
%389 = tensor.empty() : tensor<1x128x64xf32>
%390 = linalg.batch_matmul ins(%387, %388 : tensor<1x128x32xf32>, tensor<1x32x64xf32>) outs(%389 : tensor<1x128x64xf32>) -> tensor<1x128x64xf32>
%391 = tensor.empty() : tensor<1x64x128xf32>
%392 = linalg.transpose ins(%263 : tensor<1x128x64xf32>) outs(%391 : tensor<1x64x128xf32>) permutation = [0, 2, 1]
%393 = arith.constant dense<1.000000e+00> : tensor<1x32x64xf32>
%394 = tensor.empty() : tensor<1x32x128xf32>
%395 = linalg.batch_matmul ins(%393, %392 : tensor<1x32x64xf32>, tensor<1x64x128xf32>) outs(%394 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%396 = tensor.empty() : tensor<1x32x128xf32>
%397 = linalg.mul ins(%395, %378 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%396 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%398 = tensor.empty() : tensor<1x32x128xf32>
%399 = linalg.mul ins(%395, %371 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%398 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%400 = arith.constant dense<1.000000e+00> : tensor<1x32x128xf32>
%401 = tensor.empty() : tensor<1x32x128xf32>
%402 = linalg.sub ins(%400, %378 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%401 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%403 = tensor.empty() : tensor<1x32x128xf32>
%404 = linalg.mul ins(%378, %402 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%403 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%405 = tensor.empty() : tensor<1x32x128xf32>
%406 = linalg.mul ins(%399, %404 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%405 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%407 = tensor.empty() : tensor<1x64x32xf32>
%408 = linalg.transpose ins(%369 : tensor<1x32x64xf32>) outs(%407 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%409 = tensor.empty() : tensor<1x32x128xf32>
%410 = linalg.add ins(%397, %406 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%409 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%411 = tensor.empty() : tensor<1x64x128xf32>
%412 = linalg.batch_matmul ins(%408, %410 : tensor<1x64x32xf32>, tensor<1x32x128xf32>) outs(%411 : tensor<1x64x128xf32>) -> tensor<1x64x128xf32>
%413 = tensor.empty() : tensor<1x128x64xf32>
%414 = linalg.transpose ins(%262 : tensor<1x64x128xf32>) outs(%413 : tensor<1x128x64xf32>) permutation = [0, 2, 1]
%415 = tensor.empty() : tensor<1x32x128xf32>
%416 = linalg.add ins(%397, %406 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%415 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%417 = tensor.empty() : tensor<1x32x64xf32>
%418 = linalg.batch_matmul ins(%416, %414 : tensor<1x32x128xf32>, tensor<1x128x64xf32>) outs(%417 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%419 = tensor.empty() : tensor<1x32x64xf32>
%420 = linalg.mul ins(%418, %367 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%419 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%421 = tensor.empty() : tensor<1x32x64xf32>
%422 = linalg.mul ins(%418, %365 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%421 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%425 = arith.constant 0.000000e+00 : f32
%423 = tensor.empty() : tensor<1x1x64xf32>
%424 = linalg.fill ins(%425 : f32) outs(%423 : tensor<1x1x64xf32>) -> tensor<1x1x64xf32>
%426 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, 0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]
} ins(%422 : tensor<1x32x64xf32>) outs(%424 : tensor<1x1x64xf32>) {
^bb0(%a0: f32, %a1: f32):
%427 = arith.addf %a0, %a1 : f32
linalg.yield %427 : f32
} -> tensor<1x1x64xf32>
%428 = tensor.empty() : tensor<1x32x64xf32>
%429 = linalg.mul ins(%420, %363 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%428 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%430 = tensor.empty() : tensor<1x32x64xf32>
%431 = linalg.mul ins(%420, %345 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%430 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%434 = arith.constant 0.000000e+00 : f32
%432 = tensor.empty() : tensor<1x32x1xf32>
%433 = linalg.fill ins(%434 : f32) outs(%432 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%435 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%431 : tensor<1x32x64xf32>) outs(%433 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%436 = arith.addf %a0, %a1 : f32
linalg.yield %436 : f32
} -> tensor<1x32x1xf32>
%437 = arith.constant dense<-5.000000e-01> : tensor<1x32x1xf32>
%438 = tensor.empty() : tensor<1x32x1xf32>
%439 = linalg.mul ins(%361, %361 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%438 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%440 = tensor.empty() : tensor<1x32x1xf32>
%441 = linalg.mul ins(%439, %361 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%440 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%442 = tensor.empty() : tensor<1x32x1xf32>
%443 = linalg.mul ins(%437, %441 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%442 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%444 = tensor.empty() : tensor<1x32x1xf32>
%445 = linalg.mul ins(%435, %443 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%444 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%446 = tensor.empty() : tensor<1x32x1xf32>
%447 = linalg.mul ins(%445, %273 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%446 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%448 = tensor.empty() : tensor<1x32x1xf32>
%449 = linalg.mul ins(%445, %351 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%448 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%450 = tensor.empty() : tensor<1x32x64xf32>
%451 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%447 : tensor<1x32x1xf32>) outs(%450 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%452 = tensor.empty() : tensor<1x32x64xf32>
%453 = linalg.mul ins(%451, %345 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%452 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%454 = tensor.empty() : tensor<1x32x64xf32>
%455 = linalg.mul ins(%451, %345 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%454 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%456 = tensor.empty() : tensor<1x64x32xf32>
%457 = linalg.transpose ins(%341 : tensor<1x32x64xf32>) outs(%456 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%458 = arith.constant dense<1.000000e+00> : tensor<1x32x64xf32>
%459 = tensor.empty() : tensor<1x32x64xf32>
%460 = linalg.add ins(%458, %429 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%459 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%461 = tensor.empty() : tensor<1x32x64xf32>
%462 = linalg.add ins(%460, %453 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%461 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%463 = tensor.empty() : tensor<1x32x64xf32>
%464 = linalg.add ins(%462, %455 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%463 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%465 = tensor.empty() : tensor<1x64x64xf32>
%466 = linalg.batch_matmul ins(%457, %464 : tensor<1x64x32xf32>, tensor<1x32x64xf32>) outs(%465 : tensor<1x64x64xf32>) -> tensor<1x64x64xf32>
%467 = tensor.empty() : tensor<1x64x64xf32>
%468 = linalg.transpose ins(%260 : tensor<1x64x64xf32>) outs(%467 : tensor<1x64x64xf32>) permutation = [0, 2, 1]
%469 = arith.constant dense<1.000000e+00> : tensor<1x32x64xf32>
%470 = tensor.empty() : tensor<1x32x64xf32>
%471 = linalg.add ins(%469, %429 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%470 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%472 = tensor.empty() : tensor<1x32x64xf32>
%473 = linalg.add ins(%471, %453 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%472 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%474 = tensor.empty() : tensor<1x32x64xf32>
%475 = linalg.add ins(%473, %455 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%474 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%476 = tensor.empty() : tensor<1x32x64xf32>
%477 = linalg.batch_matmul ins(%475, %468 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%476 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%478 = tensor.empty() : tensor<1x32x32xf32>
%479 = linalg.transpose ins(%339 : tensor<1x32x32xf32>) outs(%478 : tensor<1x32x32xf32>) permutation = [0, 2, 1]
%480 = tensor.empty() : tensor<1x32x64xf32>
%481 = linalg.batch_matmul ins(%479, %477 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%480 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%482 = tensor.empty() : tensor<1x64x32xf32>
%483 = linalg.transpose ins(%299 : tensor<1x32x64xf32>) outs(%482 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%484 = tensor.empty() : tensor<1x32x32xf32>
%485 = linalg.batch_matmul ins(%477, %483 : tensor<1x32x64xf32>, tensor<1x64x32xf32>) outs(%484 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%486 = tensor.empty() : tensor<1x32x32xf32>
%487 = linalg.mul ins(%339, %485 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%486 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%490 = arith.constant 0.000000e+00 : f32
%488 = tensor.empty() : tensor<1x32x1xf32>
%489 = linalg.fill ins(%490 : f32) outs(%488 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%491 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%487 : tensor<1x32x32xf32>) outs(%489 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%492 = arith.addf %a0, %a1 : f32
linalg.yield %492 : f32
} -> tensor<1x32x1xf32>
%493 = tensor.empty() : tensor<1x32x32xf32>
%494 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%491 : tensor<1x32x1xf32>) outs(%493 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%495 = tensor.empty() : tensor<1x32x32xf32>
%496 = linalg.sub ins(%485, %494 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%495 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%497 = tensor.empty() : tensor<1x32x32xf32>
%498 = linalg.mul ins(%339, %496 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%497 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%499 = tensor.empty() : tensor<1x32x32xf32>
%500 = linalg.mul ins(%498, %306 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%499 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%501 = tensor.empty() : tensor<1x32x32xf32>
%502 = linalg.mul ins(%498, %303 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%501 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%505 = arith.constant 0.000000e+00 : f32
%503 = tensor.empty() : tensor<1x32x1xf32>
%504 = linalg.fill ins(%505 : f32) outs(%503 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%506 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%502 : tensor<1x32x32xf32>) outs(%504 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%507 = arith.addf %a0, %a1 : f32
linalg.yield %507 : f32
} -> tensor<1x32x1xf32>
%510 = arith.constant 0.000000e+00 : f32
%508 = tensor.empty() : tensor<1x1x1xf32>
%509 = linalg.fill ins(%510 : f32) outs(%508 : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
%511 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, 0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]
} ins(%506 : tensor<1x32x1xf32>) outs(%509 : tensor<1x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%512 = arith.addf %a0, %a1 : f32
linalg.yield %512 : f32
} -> tensor<1x1x1xf32>
%513 = tensor.empty() : tensor<1x64x32xf32>
%514 = linalg.transpose ins(%295 : tensor<1x32x64xf32>) outs(%513 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%515 = tensor.empty() : tensor<1x64x32xf32>
%516 = linalg.batch_matmul ins(%514, %500 : tensor<1x64x32xf32>, tensor<1x32x32xf32>) outs(%515 : tensor<1x64x32xf32>) -> tensor<1x64x32xf32>
%517 = tensor.empty() : tensor<1x32x64xf32>
%518 = linalg.transpose ins(%301 : tensor<1x64x32xf32>) outs(%517 : tensor<1x32x64xf32>) permutation = [0, 2, 1]
%519 = tensor.empty() : tensor<1x32x64xf32>
%520 = linalg.batch_matmul ins(%500, %518 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%519 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%521 = tensor.empty() : tensor<1x32x64xf32>
%522 = linalg.transpose ins(%516 : tensor<1x64x32xf32>) outs(%521 : tensor<1x32x64xf32>) permutation = [0, 2, 1]
%523 = tensor.empty() : tensor<1x64x32xf32>
%524 = linalg.transpose ins(%293 : tensor<1x32x64xf32>) outs(%523 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%525 = tensor.empty() : tensor<1x64x64xf32>
%526 = linalg.batch_matmul ins(%524, %481 : tensor<1x64x32xf32>, tensor<1x32x64xf32>) outs(%525 : tensor<1x64x64xf32>) -> tensor<1x64x64xf32>
%527 = tensor.empty() : tensor<1x64x64xf32>
%528 = linalg.transpose ins(%259 : tensor<1x64x64xf32>) outs(%527 : tensor<1x64x64xf32>) permutation = [0, 2, 1]
%529 = tensor.empty() : tensor<1x32x64xf32>
%530 = linalg.batch_matmul ins(%481, %528 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%529 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%531 = tensor.empty() : tensor<1x64x32xf32>
%532 = linalg.transpose ins(%293 : tensor<1x32x64xf32>) outs(%531 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%533 = tensor.empty() : tensor<1x64x64xf32>
%534 = linalg.batch_matmul ins(%532, %522 : tensor<1x64x32xf32>, tensor<1x32x64xf32>) outs(%533 : tensor<1x64x64xf32>) -> tensor<1x64x64xf32>
%535 = tensor.empty() : tensor<1x64x64xf32>
%536 = linalg.transpose ins(%258 : tensor<1x64x64xf32>) outs(%535 : tensor<1x64x64xf32>) permutation = [0, 2, 1]
%537 = tensor.empty() : tensor<1x32x64xf32>
%538 = linalg.batch_matmul ins(%522, %536 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%537 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%539 = tensor.empty() : tensor<1x64x32xf32>
%540 = linalg.transpose ins(%293 : tensor<1x32x64xf32>) outs(%539 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%541 = tensor.empty() : tensor<1x64x64xf32>
%542 = linalg.batch_matmul ins(%540, %520 : tensor<1x64x32xf32>, tensor<1x32x64xf32>) outs(%541 : tensor<1x64x64xf32>) -> tensor<1x64x64xf32>
%543 = tensor.empty() : tensor<1x64x64xf32>
%544 = linalg.transpose ins(%257 : tensor<1x64x64xf32>) outs(%543 : tensor<1x64x64xf32>) permutation = [0, 2, 1]
%545 = tensor.empty() : tensor<1x32x64xf32>
%546 = linalg.batch_matmul ins(%520, %544 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%545 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%547 = tensor.empty() : tensor<1x32x64xf32>
%548 = linalg.add ins(%530, %538 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%547 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%549 = tensor.empty() : tensor<1x32x64xf32>
%550 = linalg.add ins(%548, %546 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%549 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%551 = tensor.empty() : tensor<1x32x64xf32>
%552 = linalg.mul ins(%550, %291 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%551 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%553 = tensor.empty() : tensor<1x32x64xf32>
%554 = linalg.add ins(%530, %538 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%553 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%555 = tensor.empty() : tensor<1x32x64xf32>
%556 = linalg.add ins(%554, %546 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%555 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%557 = tensor.empty() : tensor<1x32x64xf32>
%558 = linalg.mul ins(%556, %289 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%557 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%561 = arith.constant 0.000000e+00 : f32
%559 = tensor.empty() : tensor<1x1x64xf32>
%560 = linalg.fill ins(%561 : f32) outs(%559 : tensor<1x1x64xf32>) -> tensor<1x1x64xf32>
%562 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, 0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]
} ins(%558 : tensor<1x32x64xf32>) outs(%560 : tensor<1x1x64xf32>) {
^bb0(%a0: f32, %a1: f32):
%563 = arith.addf %a0, %a1 : f32
linalg.yield %563 : f32
} -> tensor<1x1x64xf32>
%564 = tensor.empty() : tensor<1x32x64xf32>
%565 = linalg.mul ins(%552, %287 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%564 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%566 = tensor.empty() : tensor<1x32x64xf32>
%567 = linalg.mul ins(%552, %255 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%566 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%570 = arith.constant 0.000000e+00 : f32
%568 = tensor.empty() : tensor<1x32x1xf32>
%569 = linalg.fill ins(%570 : f32) outs(%568 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%571 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%567 : tensor<1x32x64xf32>) outs(%569 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%572 = arith.addf %a0, %a1 : f32
linalg.yield %572 : f32
} -> tensor<1x32x1xf32>
%573 = arith.constant dense<-5.000000e-01> : tensor<1x32x1xf32>
%574 = tensor.empty() : tensor<1x32x1xf32>
%575 = linalg.mul ins(%285, %285 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%574 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%576 = tensor.empty() : tensor<1x32x1xf32>
%577 = linalg.mul ins(%575, %285 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%576 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%578 = tensor.empty() : tensor<1x32x1xf32>
%579 = linalg.mul ins(%573, %577 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%578 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%580 = tensor.empty() : tensor<1x32x1xf32>
%581 = linalg.mul ins(%571, %579 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%580 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%582 = tensor.empty() : tensor<1x32x1xf32>
%583 = linalg.add ins(%445, %581 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%582 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%586 = arith.constant 0.000000e+00 : f32
%584 = tensor.empty() : tensor<1x1x1xf32>
%585 = linalg.fill ins(%586 : f32) outs(%584 : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
%587 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, 0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]
} ins(%583 : tensor<1x32x1xf32>) outs(%585 : tensor<1x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%588 = arith.addf %a0, %a1 : f32
linalg.yield %588 : f32
} -> tensor<1x1x1xf32>
%589 = tensor.empty() : tensor<1x32x1xf32>
%590 = linalg.mul ins(%581, %273 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%589 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%591 = tensor.empty() : tensor<1x32x1xf32>
%592 = linalg.mul ins(%581, %269 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%591 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%593 = tensor.empty() : tensor<1x32x1xf32>
%594 = linalg.add ins(%449, %592 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%593 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%597 = arith.constant 0.000000e+00 : f32
%595 = tensor.empty() : tensor<1x1x1xf32>
%596 = linalg.fill ins(%597 : f32) outs(%595 : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
%598 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, 0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]
} ins(%594 : tensor<1x32x1xf32>) outs(%596 : tensor<1x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%599 = arith.addf %a0, %a1 : f32
linalg.yield %599 : f32
} -> tensor<1x1x1xf32>
%600 = tensor.empty() : tensor<1x32x64xf32>
%601 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%590 : tensor<1x32x1xf32>) outs(%600 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%602 = tensor.empty() : tensor<1x32x64xf32>
%603 = linalg.mul ins(%601, %255 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%602 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%604 = tensor.empty() : tensor<1x32x64xf32>
%605 = linalg.mul ins(%601, %255 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%604 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%606 = arith.constant dense<1.000000e+00> : tensor<1x32x64xf32>
%607 = tensor.empty() : tensor<1x32x64xf32>
%608 = linalg.add ins(%606, %429 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%607 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%609 = tensor.empty() : tensor<1x32x64xf32>
%610 = linalg.add ins(%608, %453 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%609 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%611 = tensor.empty() : tensor<1x32x64xf32>
%612 = linalg.add ins(%610, %455 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%611 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%613 = tensor.empty() : tensor<1x32x64xf32>
%614 = linalg.add ins(%612, %565 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%613 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%615 = tensor.empty() : tensor<1x32x64xf32>
%616 = linalg.add ins(%614, %603 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%615 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%617 = tensor.empty() : tensor<1x32x64xf32>
%618 = linalg.add ins(%616, %605 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%617 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return
}
}