Example
storylm0_adamw_step.kina
let adamw_e = fn(
w: tensor<F32, [8, 4], @host>,
g: tensor<F32, [8, 4], @host>,
m: tensor<F32, [8, 4], @host>,
v: tensor<F32, [8, 4], @host>,
lr: tensor<F32, [1, 1], @host>,
beta1: tensor<F32, [1, 1], @host>,
beta2: tensor<F32, [1, 1], @host>,
eps: tensor<F32, [1, 1], @host>,
wd: tensor<F32, [1, 1], @host>,
beta1_pow: tensor<F32, [1, 1], @host>,
beta2_pow: tensor<F32, [1, 1], @host>,
one: tensor<F32, [1, 1], @host>
) -> (
tensor<F32, [8, 4], @host>,
tensor<F32, [8, 4], @host>,
tensor<F32, [8, 4], @host>
) {
let one_minus_beta1 = one - beta1 in
let one_minus_beta2 = one - beta2 in
let beta1_b = beta1.broadcast([8, 4]) in
let beta2_b = beta2.broadcast([8, 4]) in
let one_minus_beta1_b = one_minus_beta1.broadcast([8, 4]) in
let one_minus_beta2_b = one_minus_beta2.broadcast([8, 4]) in
let m_keep = beta1_b * m in
let m_grad = one_minus_beta1_b * g in
let m_next = m_keep + m_grad in
let g2 = g * g in
let v_keep = beta2_b * v in
let v_grad = one_minus_beta2_b * g2 in
let v_next = v_keep + v_grad in
let bc1 = one - beta1_pow in
let bc2 = one - beta2_pow in
let v_sqrt = sqrt(v_next) in
let bc2_sqrt = sqrt(bc2) in
let bc2_sqrt_b = bc2_sqrt.broadcast([8, 4]) in
let v_norm = v_sqrt / bc2_sqrt_b in
let eps_b = eps.broadcast([8, 4]) in
let denom = v_norm + eps_b in
let step_size = lr / bc1 in
let step_size_b = step_size.broadcast([8, 4]) in
let m_over_denom = m_next / denom in
let adam_step = step_size_b * m_over_denom in
let decay = one - lr * wd in
let decay_b = decay.broadcast([8, 4]) in
let w_decayed = w * decay_b in
let w_next = w_decayed - adam_step in
(w_next, m_next, v_next)
} in
let adamw_w = fn(
w: tensor<F32, [4, 4], @host>,
g: tensor<F32, [4, 4], @host>,
m: tensor<F32, [4, 4], @host>,
v: tensor<F32, [4, 4], @host>,
lr: tensor<F32, [1, 1], @host>,
beta1: tensor<F32, [1, 1], @host>,
beta2: tensor<F32, [1, 1], @host>,
eps: tensor<F32, [1, 1], @host>,
wd: tensor<F32, [1, 1], @host>,
beta1_pow: tensor<F32, [1, 1], @host>,
beta2_pow: tensor<F32, [1, 1], @host>,
one: tensor<F32, [1, 1], @host>
) -> (
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>
) {
let one_minus_beta1 = one - beta1 in
let one_minus_beta2 = one - beta2 in
let beta1_b = beta1.broadcast([4, 4]) in
let beta2_b = beta2.broadcast([4, 4]) in
let one_minus_beta1_b = one_minus_beta1.broadcast([4, 4]) in
let one_minus_beta2_b = one_minus_beta2.broadcast([4, 4]) in
let m_keep = beta1_b * m in
let m_grad = one_minus_beta1_b * g in
let m_next = m_keep + m_grad in
let g2 = g * g in
let v_keep = beta2_b * v in
let v_grad = one_minus_beta2_b * g2 in
let v_next = v_keep + v_grad in
let bc1 = one - beta1_pow in
let bc2 = one - beta2_pow in
let v_sqrt = sqrt(v_next) in
let bc2_sqrt = sqrt(bc2) in
let bc2_sqrt_b = bc2_sqrt.broadcast([4, 4]) in
let v_norm = v_sqrt / bc2_sqrt_b in
let eps_b = eps.broadcast([4, 4]) in
let denom = v_norm + eps_b in
let step_size = lr / bc1 in
let step_size_b = step_size.broadcast([4, 4]) in
let m_over_denom = m_next / denom in
let adam_step = step_size_b * m_over_denom in
let decay = one - lr * wd in
let decay_b = decay.broadcast([4, 4]) in
let w_decayed = w * decay_b in
let w_next = w_decayed - adam_step in
(w_next, m_next, v_next)
} in
let adamw_lm = fn(
w: tensor<F32, [4, 8], @host>,
g: tensor<F32, [4, 8], @host>,
m: tensor<F32, [4, 8], @host>,
v: tensor<F32, [4, 8], @host>,
lr: tensor<F32, [1, 1], @host>,
beta1: tensor<F32, [1, 1], @host>,
beta2: tensor<F32, [1, 1], @host>,
eps: tensor<F32, [1, 1], @host>,
wd: tensor<F32, [1, 1], @host>,
beta1_pow: tensor<F32, [1, 1], @host>,
beta2_pow: tensor<F32, [1, 1], @host>,
one: tensor<F32, [1, 1], @host>
) -> (
tensor<F32, [4, 8], @host>,
tensor<F32, [4, 8], @host>,
tensor<F32, [4, 8], @host>
) {
let one_minus_beta1 = one - beta1 in
let one_minus_beta2 = one - beta2 in
let beta1_b = beta1.broadcast([4, 8]) in
let beta2_b = beta2.broadcast([4, 8]) in
let one_minus_beta1_b = one_minus_beta1.broadcast([4, 8]) in
let one_minus_beta2_b = one_minus_beta2.broadcast([4, 8]) in
let m_keep = beta1_b * m in
let m_grad = one_minus_beta1_b * g in
let m_next = m_keep + m_grad in
let g2 = g * g in
let v_keep = beta2_b * v in
let v_grad = one_minus_beta2_b * g2 in
let v_next = v_keep + v_grad in
let bc1 = one - beta1_pow in
let bc2 = one - beta2_pow in
let v_sqrt = sqrt(v_next) in
let bc2_sqrt = sqrt(bc2) in
let bc2_sqrt_b = bc2_sqrt.broadcast([4, 8]) in
let v_norm = v_sqrt / bc2_sqrt_b in
let eps_b = eps.broadcast([4, 8]) in
let denom = v_norm + eps_b in
let step_size = lr / bc1 in
let step_size_b = step_size.broadcast([4, 8]) in
let m_over_denom = m_next / denom in
let adam_step = step_size_b * m_over_denom in
let decay = one - lr * wd in
let decay_b = decay.broadcast([4, 8]) in
let w_decayed = w * decay_b in
let w_next = w_decayed - adam_step in
(w_next, m_next, v_next)
} in
let storylm0_nll = fn(
token_ids: tensor<I64, [1, 4], @host>,
targets: tensor<I64, [1, 4], @host>,
token_embedding: tensor<F32, [8, 4], @host>,
causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
wq: tensor<F32, [4, 4], @host>,
wk: tensor<F32, [4, 4], @host>,
wv: tensor<F32, [4, 4], @host>,
wo: tensor<F32, [4, 4], @host>,
lm_head: tensor<F32, [4, 8], @host>
) -> tensor<F32, [1, 4], @host> {
let x = token_embedding.gather(token_ids) in
// Unit batch dim is a volume-preserving reshape. Rank-changing broadcast
// still has no VJP in the restricted pass.
let wq_broadcast = wq.reshape([1, 4, 4]) in
let wk_broadcast = wk.reshape([1, 4, 4]) in
let wv_broadcast = wv.reshape([1, 4, 4]) in
let wo_broadcast = wo.reshape([1, 4, 4]) in
let lm_head_broadcast = lm_head.reshape([1, 4, 8]) in
let q_flat = x @ wq_broadcast in
let k_flat = x @ wk_broadcast in
let v_flat = x @ wv_broadcast in
let q_unsplit = q_flat.reshape([1, 4, 2, 2]) in
let q = q_unsplit.transpose([0, 2, 1, 3]) in
let k_unsplit = k_flat.reshape([1, 4, 2, 2]) in
let k_t = k_unsplit.transpose([0, 2, 3, 1]) in
let v_unsplit = v_flat.reshape([1, 4, 2, 2]) in
let v = v_unsplit.transpose([0, 2, 1, 3]) in
let scores = q @ k_t in
let scale_broadcast = attention_scale.broadcast([1, 2, 4, 4]) in
let mask_broadcast = causal_mask.broadcast([1, 2, 4, 4]) in
let scaled_scores = scores * scale_broadcast in
let masked_scores = scaled_scores + mask_broadcast in
let probs = masked_scores.softmax(axis=3) in
let context = probs @ v in
let context_t = context.transpose([0, 2, 1, 3]) in
let merged = context_t.reshape([1, 4, 4]) in
let projected = merged @ wo_broadcast in
let residual = x + projected in
let logits = residual @ lm_head_broadcast in
let log_probs = logits.log_softmax(axis=2) in
let picked = log_probs.take_along_axis(targets, axis=2) in
let nll = -picked in
nll
} in
let step = fn(
token_ids: tensor<I64, [1, 4], @host>,
targets: tensor<I64, [1, 4], @host>,
token_embedding: tensor<F32, [8, 4], @host>,
causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
wq: tensor<F32, [4, 4], @host>,
wk: tensor<F32, [4, 4], @host>,
wv: tensor<F32, [4, 4], @host>,
wo: tensor<F32, [4, 4], @host>,
lm_head: tensor<F32, [4, 8], @host>,
m_emb: tensor<F32, [8, 4], @host>,
v_emb: tensor<F32, [8, 4], @host>,
m_wq: tensor<F32, [4, 4], @host>,
v_wq: tensor<F32, [4, 4], @host>,
m_wk: tensor<F32, [4, 4], @host>,
v_wk: tensor<F32, [4, 4], @host>,
m_wv: tensor<F32, [4, 4], @host>,
v_wv: tensor<F32, [4, 4], @host>,
m_wo: tensor<F32, [4, 4], @host>,
v_wo: tensor<F32, [4, 4], @host>,
m_lm: tensor<F32, [4, 8], @host>,
v_lm: tensor<F32, [4, 8], @host>,
lr: tensor<F32, [1, 1], @host>,
beta1: tensor<F32, [1, 1], @host>,
beta2: tensor<F32, [1, 1], @host>,
eps: tensor<F32, [1, 1], @host>,
wd: tensor<F32, [1, 1], @host>,
beta1_pow: tensor<F32, [1, 1], @host>,
beta2_pow: tensor<F32, [1, 1], @host>,
one: tensor<F32, [1, 1], @host>
) -> (
tensor<F32, [1, 4], @host>,
tensor<F32, [8, 4], @host>,
tensor<F32, [8, 4], @host>,
tensor<F32, [8, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 4], @host>,
tensor<F32, [4, 8], @host>,
tensor<F32, [4, 8], @host>,
tensor<F32, [4, 8], @host>
) {
let (nll, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm) =
handle storylm0_nll(
token_ids, targets, token_embedding, causal_mask, attention_scale,
wq, wk, wv, wo, lm_head
) with reverseMode in
let (emb_next, m_emb_next, v_emb_next) =
adamw_e(
token_embedding, g_emb, m_emb, v_emb, lr, beta1, beta2, eps, wd,
beta1_pow, beta2_pow, one
) in
let (wq_next, m_wq_next, v_wq_next) =
adamw_w(wq, g_wq, m_wq, v_wq, lr, beta1, beta2, eps, wd, beta1_pow, beta2_pow, one) in
let (wk_next, m_wk_next, v_wk_next) =
adamw_w(wk, g_wk, m_wk, v_wk, lr, beta1, beta2, eps, wd, beta1_pow, beta2_pow, one) in
let (wv_next, m_wv_next, v_wv_next) =
adamw_w(wv, g_wv, m_wv, v_wv, lr, beta1, beta2, eps, wd, beta1_pow, beta2_pow, one) in
let (wo_next, m_wo_next, v_wo_next) =
adamw_w(wo, g_wo, m_wo, v_wo, lr, beta1, beta2, eps, wd, beta1_pow, beta2_pow, one) in
let (lm_next, m_lm_next, v_lm_next) =
adamw_lm(
lm_head, g_lm, m_lm, v_lm, lr, beta1, beta2, eps, wd, beta1_pow,
beta2_pow, one
) in
(
nll,
emb_next, m_emb_next, v_emb_next,
wq_next, m_wq_next, v_wq_next,
wk_next, m_wk_next, v_wk_next,
wv_next, m_wv_next, v_wv_next,
wo_next, m_wo_next, v_wo_next,
lm_next, m_lm_next, v_lm_next
)
} in
let token_ids = tensor<I64, [1, 4], @host> in
let targets = tensor<I64, [1, 4], @host> in
let token_embedding = tensor<F32, [8, 4], @host> in
let causal_mask = tensor<F32, [1, 1, 4, 4], @host> in
let attention_scale = tensor<F32, [1, 1, 1, 1], @host> in
let wq = tensor<F32, [4, 4], @host> in
let wk = tensor<F32, [4, 4], @host> in
let wv = tensor<F32, [4, 4], @host> in
let wo = tensor<F32, [4, 4], @host> in
let lm_head = tensor<F32, [4, 8], @host> in
let m_emb = tensor<F32, [8, 4], @host> in
let v_emb = tensor<F32, [8, 4], @host> in
let m_wq = tensor<F32, [4, 4], @host> in
let v_wq = tensor<F32, [4, 4], @host> in
let m_wk = tensor<F32, [4, 4], @host> in
let v_wk = tensor<F32, [4, 4], @host> in
let m_wv = tensor<F32, [4, 4], @host> in
let v_wv = tensor<F32, [4, 4], @host> in
let m_wo = tensor<F32, [4, 4], @host> in
let v_wo = tensor<F32, [4, 4], @host> in
let m_lm = tensor<F32, [4, 8], @host> in
let v_lm = tensor<F32, [4, 8], @host> in
let lr = tensor<F32, [1, 1], @host> in
let beta1 = tensor<F32, [1, 1], @host> in
let beta2 = tensor<F32, [1, 1], @host> in
let eps = tensor<F32, [1, 1], @host> in
let wd = tensor<F32, [1, 1], @host> in
let beta1_pow = tensor<F32, [1, 1], @host> in
let beta2_pow = tensor<F32, [1, 1], @host> in
let one = tensor<F32, [1, 1], @host> in
step(
token_ids, targets, token_embedding, causal_mask, attention_scale,
wq, wk, wv, wo, lm_head,
m_emb, v_emb, m_wq, v_wq, m_wk, v_wk, m_wv, v_wv, m_wo, v_wo, m_lm, v_lm,
lr, beta1, beta2, eps, wd, beta1_pow, beta2_pow, one
)
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 = arith.constant dense<0.000000e+00> : 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 = arith.constant dense<0.000000e+00> : 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 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%57 = linalg.batch_matmul ins(%51, %52 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%56 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%58 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%59 = linalg.batch_matmul ins(%51, %53 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%58 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%60 = arith.constant dense<0.000000e+00> : 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 = arith.constant dense<0.000000e+00> : 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 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%103 = linalg.batch_matmul ins(%101, %61 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%102 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%104 = arith.constant dense<0.000000e+00> : 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<8x4xf32>, %126: tensor<8x4xf32>, %127: tensor<8x4xf32>, %128: tensor<8x4xf32>, %129: tensor<1x1xf32>, %130: tensor<1x1xf32>, %131: tensor<1x1xf32>, %132: tensor<1x1xf32>, %133: tensor<1x1xf32>, %134: tensor<1x1xf32>, %135: tensor<1x1xf32>, %136: tensor<1x1xf32>) -> (tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>) attributes {llvm.emit_c_interface} {
%137 = tensor.empty() : tensor<1x1xf32>
%138 = linalg.sub ins(%136, %130 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%137 : tensor<1x1xf32>) -> tensor<1x1xf32>
%139 = tensor.empty() : tensor<1x1xf32>
%140 = linalg.sub ins(%136, %131 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%139 : tensor<1x1xf32>) -> tensor<1x1xf32>
%141 = tensor.empty() : tensor<8x4xf32>
%142 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%130 : tensor<1x1xf32>) outs(%141 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%143 = tensor.empty() : tensor<8x4xf32>
%144 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%131 : tensor<1x1xf32>) outs(%143 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%145 = tensor.empty() : tensor<8x4xf32>
%146 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%138 : tensor<1x1xf32>) outs(%145 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%147 = tensor.empty() : tensor<8x4xf32>
%148 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%140 : tensor<1x1xf32>) outs(%147 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%149 = tensor.empty() : tensor<8x4xf32>
%150 = linalg.mul ins(%142, %127 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%149 : tensor<8x4xf32>) -> tensor<8x4xf32>
%151 = tensor.empty() : tensor<8x4xf32>
%152 = linalg.mul ins(%146, %126 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%151 : tensor<8x4xf32>) -> tensor<8x4xf32>
%153 = tensor.empty() : tensor<8x4xf32>
%154 = linalg.add ins(%150, %152 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%153 : tensor<8x4xf32>) -> tensor<8x4xf32>
%155 = tensor.empty() : tensor<8x4xf32>
%156 = linalg.mul ins(%126, %126 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%155 : tensor<8x4xf32>) -> tensor<8x4xf32>
%157 = tensor.empty() : tensor<8x4xf32>
%158 = linalg.mul ins(%144, %128 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%157 : tensor<8x4xf32>) -> tensor<8x4xf32>
%159 = tensor.empty() : tensor<8x4xf32>
%160 = linalg.mul ins(%148, %156 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%159 : tensor<8x4xf32>) -> tensor<8x4xf32>
%161 = tensor.empty() : tensor<8x4xf32>
%162 = linalg.add ins(%158, %160 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%161 : tensor<8x4xf32>) -> tensor<8x4xf32>
%163 = tensor.empty() : tensor<1x1xf32>
%164 = linalg.sub ins(%136, %134 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%163 : tensor<1x1xf32>) -> tensor<1x1xf32>
%165 = tensor.empty() : tensor<1x1xf32>
%166 = linalg.sub ins(%136, %135 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%165 : tensor<1x1xf32>) -> tensor<1x1xf32>
%168 = tensor.empty() : tensor<8x4xf32>
%169 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%162 : tensor<8x4xf32>) outs(%168 : tensor<8x4xf32>) {
^bb0(%x: f32, %_: f32):
%167 = math.sqrt %x : f32
linalg.yield %167 : f32
} -> tensor<8x4xf32>
%171 = tensor.empty() : tensor<1x1xf32>
%172 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%166 : tensor<1x1xf32>) outs(%171 : tensor<1x1xf32>) {
^bb0(%x: f32, %_: f32):
%170 = math.sqrt %x : f32
linalg.yield %170 : f32
} -> tensor<1x1xf32>
%173 = tensor.empty() : tensor<8x4xf32>
%174 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%172 : tensor<1x1xf32>) outs(%173 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%175 = tensor.empty() : tensor<8x4xf32>
%176 = linalg.div ins(%169, %174 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%175 : tensor<8x4xf32>) -> tensor<8x4xf32>
%177 = tensor.empty() : tensor<8x4xf32>
%178 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%132 : tensor<1x1xf32>) outs(%177 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%179 = tensor.empty() : tensor<8x4xf32>
%180 = linalg.add ins(%176, %178 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%179 : tensor<8x4xf32>) -> tensor<8x4xf32>
%181 = tensor.empty() : tensor<1x1xf32>
%182 = linalg.div ins(%129, %164 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%181 : tensor<1x1xf32>) -> tensor<1x1xf32>
%183 = tensor.empty() : tensor<8x4xf32>
%184 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%182 : tensor<1x1xf32>) outs(%183 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%185 = tensor.empty() : tensor<8x4xf32>
%186 = linalg.div ins(%154, %180 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%185 : tensor<8x4xf32>) -> tensor<8x4xf32>
%187 = tensor.empty() : tensor<8x4xf32>
%188 = linalg.mul ins(%184, %186 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%187 : tensor<8x4xf32>) -> tensor<8x4xf32>
%189 = tensor.empty() : tensor<1x1xf32>
%190 = linalg.mul ins(%129, %133 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%189 : tensor<1x1xf32>) -> tensor<1x1xf32>
%191 = tensor.empty() : tensor<1x1xf32>
%192 = linalg.sub ins(%136, %190 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%191 : tensor<1x1xf32>) -> tensor<1x1xf32>
%193 = tensor.empty() : tensor<8x4xf32>
%194 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%192 : tensor<1x1xf32>) outs(%193 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%195 = tensor.empty() : tensor<8x4xf32>
%196 = linalg.mul ins(%125, %194 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%195 : tensor<8x4xf32>) -> tensor<8x4xf32>
%197 = tensor.empty() : tensor<8x4xf32>
%198 = linalg.sub ins(%196, %188 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%197 : tensor<8x4xf32>) -> tensor<8x4xf32>
return %198, %154, %162 : tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>
}
func.func @fn_val_199(%200: tensor<4x4xf32>, %201: tensor<4x4xf32>, %202: tensor<4x4xf32>, %203: tensor<4x4xf32>, %204: tensor<1x1xf32>, %205: tensor<1x1xf32>, %206: tensor<1x1xf32>, %207: tensor<1x1xf32>, %208: tensor<1x1xf32>, %209: tensor<1x1xf32>, %210: tensor<1x1xf32>, %211: tensor<1x1xf32>) -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>) attributes {llvm.emit_c_interface} {
%212 = tensor.empty() : tensor<1x1xf32>
%213 = linalg.sub ins(%211, %205 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%212 : tensor<1x1xf32>) -> tensor<1x1xf32>
%214 = tensor.empty() : tensor<1x1xf32>
%215 = linalg.sub ins(%211, %206 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%214 : tensor<1x1xf32>) -> tensor<1x1xf32>
%216 = tensor.empty() : tensor<4x4xf32>
%217 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%205 : tensor<1x1xf32>) outs(%216 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%218 = tensor.empty() : tensor<4x4xf32>
%219 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%206 : tensor<1x1xf32>) outs(%218 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%220 = tensor.empty() : tensor<4x4xf32>
%221 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%213 : tensor<1x1xf32>) outs(%220 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%222 = tensor.empty() : tensor<4x4xf32>
%223 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%215 : tensor<1x1xf32>) outs(%222 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%224 = tensor.empty() : tensor<4x4xf32>
%225 = linalg.mul ins(%217, %202 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%224 : tensor<4x4xf32>) -> tensor<4x4xf32>
%226 = tensor.empty() : tensor<4x4xf32>
%227 = linalg.mul ins(%221, %201 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%226 : tensor<4x4xf32>) -> tensor<4x4xf32>
%228 = tensor.empty() : tensor<4x4xf32>
%229 = linalg.add ins(%225, %227 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%228 : tensor<4x4xf32>) -> tensor<4x4xf32>
%230 = tensor.empty() : tensor<4x4xf32>
%231 = linalg.mul ins(%201, %201 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%230 : tensor<4x4xf32>) -> tensor<4x4xf32>
%232 = tensor.empty() : tensor<4x4xf32>
%233 = linalg.mul ins(%219, %203 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%232 : tensor<4x4xf32>) -> tensor<4x4xf32>
%234 = tensor.empty() : tensor<4x4xf32>
%235 = linalg.mul ins(%223, %231 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%234 : tensor<4x4xf32>) -> tensor<4x4xf32>
%236 = tensor.empty() : tensor<4x4xf32>
%237 = linalg.add ins(%233, %235 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%236 : tensor<4x4xf32>) -> tensor<4x4xf32>
%238 = tensor.empty() : tensor<1x1xf32>
%239 = linalg.sub ins(%211, %209 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%238 : tensor<1x1xf32>) -> tensor<1x1xf32>
%240 = tensor.empty() : tensor<1x1xf32>
%241 = linalg.sub ins(%211, %210 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%240 : tensor<1x1xf32>) -> tensor<1x1xf32>
%243 = tensor.empty() : tensor<4x4xf32>
%244 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%237 : tensor<4x4xf32>) outs(%243 : tensor<4x4xf32>) {
^bb0(%x: f32, %_: f32):
%242 = math.sqrt %x : f32
linalg.yield %242 : f32
} -> tensor<4x4xf32>
%246 = tensor.empty() : tensor<1x1xf32>
%247 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%241 : tensor<1x1xf32>) outs(%246 : tensor<1x1xf32>) {
^bb0(%x: f32, %_: f32):
%245 = math.sqrt %x : f32
linalg.yield %245 : f32
} -> tensor<1x1xf32>
%248 = tensor.empty() : tensor<4x4xf32>
%249 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%247 : tensor<1x1xf32>) outs(%248 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%250 = tensor.empty() : tensor<4x4xf32>
%251 = linalg.div ins(%244, %249 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%250 : tensor<4x4xf32>) -> tensor<4x4xf32>
%252 = tensor.empty() : tensor<4x4xf32>
%253 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%207 : tensor<1x1xf32>) outs(%252 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%254 = tensor.empty() : tensor<4x4xf32>
%255 = linalg.add ins(%251, %253 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%254 : tensor<4x4xf32>) -> tensor<4x4xf32>
%256 = tensor.empty() : tensor<1x1xf32>
%257 = linalg.div ins(%204, %239 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%256 : tensor<1x1xf32>) -> tensor<1x1xf32>
%258 = tensor.empty() : tensor<4x4xf32>
%259 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%257 : tensor<1x1xf32>) outs(%258 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%260 = tensor.empty() : tensor<4x4xf32>
%261 = linalg.div ins(%229, %255 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%260 : tensor<4x4xf32>) -> tensor<4x4xf32>
%262 = tensor.empty() : tensor<4x4xf32>
%263 = linalg.mul ins(%259, %261 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%262 : tensor<4x4xf32>) -> tensor<4x4xf32>
%264 = tensor.empty() : tensor<1x1xf32>
%265 = linalg.mul ins(%204, %208 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%264 : tensor<1x1xf32>) -> tensor<1x1xf32>
%266 = tensor.empty() : tensor<1x1xf32>
%267 = linalg.sub ins(%211, %265 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%266 : tensor<1x1xf32>) -> tensor<1x1xf32>
%268 = tensor.empty() : tensor<4x4xf32>
%269 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%267 : tensor<1x1xf32>) outs(%268 : tensor<4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x4xf32>
%270 = tensor.empty() : tensor<4x4xf32>
%271 = linalg.mul ins(%200, %269 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%270 : tensor<4x4xf32>) -> tensor<4x4xf32>
%272 = tensor.empty() : tensor<4x4xf32>
%273 = linalg.sub ins(%271, %263 : tensor<4x4xf32>, tensor<4x4xf32>) outs(%272 : tensor<4x4xf32>) -> tensor<4x4xf32>
return %273, %229, %237 : tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>
}
func.func @fn_val_274(%275: tensor<4x8xf32>, %276: tensor<4x8xf32>, %277: tensor<4x8xf32>, %278: tensor<4x8xf32>, %279: tensor<1x1xf32>, %280: tensor<1x1xf32>, %281: tensor<1x1xf32>, %282: tensor<1x1xf32>, %283: tensor<1x1xf32>, %284: tensor<1x1xf32>, %285: tensor<1x1xf32>, %286: tensor<1x1xf32>) -> (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>) attributes {llvm.emit_c_interface} {
%287 = tensor.empty() : tensor<1x1xf32>
%288 = linalg.sub ins(%286, %280 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%287 : tensor<1x1xf32>) -> tensor<1x1xf32>
%289 = tensor.empty() : tensor<1x1xf32>
%290 = linalg.sub ins(%286, %281 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%289 : tensor<1x1xf32>) -> tensor<1x1xf32>
%291 = tensor.empty() : tensor<4x8xf32>
%292 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%280 : tensor<1x1xf32>) outs(%291 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%293 = tensor.empty() : tensor<4x8xf32>
%294 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%281 : tensor<1x1xf32>) outs(%293 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%295 = tensor.empty() : tensor<4x8xf32>
%296 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%288 : tensor<1x1xf32>) outs(%295 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%297 = tensor.empty() : tensor<4x8xf32>
%298 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%290 : tensor<1x1xf32>) outs(%297 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%299 = tensor.empty() : tensor<4x8xf32>
%300 = linalg.mul ins(%292, %277 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%299 : tensor<4x8xf32>) -> tensor<4x8xf32>
%301 = tensor.empty() : tensor<4x8xf32>
%302 = linalg.mul ins(%296, %276 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%301 : tensor<4x8xf32>) -> tensor<4x8xf32>
%303 = tensor.empty() : tensor<4x8xf32>
%304 = linalg.add ins(%300, %302 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%303 : tensor<4x8xf32>) -> tensor<4x8xf32>
%305 = tensor.empty() : tensor<4x8xf32>
%306 = linalg.mul ins(%276, %276 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%305 : tensor<4x8xf32>) -> tensor<4x8xf32>
%307 = tensor.empty() : tensor<4x8xf32>
%308 = linalg.mul ins(%294, %278 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%307 : tensor<4x8xf32>) -> tensor<4x8xf32>
%309 = tensor.empty() : tensor<4x8xf32>
%310 = linalg.mul ins(%298, %306 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%309 : tensor<4x8xf32>) -> tensor<4x8xf32>
%311 = tensor.empty() : tensor<4x8xf32>
%312 = linalg.add ins(%308, %310 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%311 : tensor<4x8xf32>) -> tensor<4x8xf32>
%313 = tensor.empty() : tensor<1x1xf32>
%314 = linalg.sub ins(%286, %284 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%313 : tensor<1x1xf32>) -> tensor<1x1xf32>
%315 = tensor.empty() : tensor<1x1xf32>
%316 = linalg.sub ins(%286, %285 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%315 : tensor<1x1xf32>) -> tensor<1x1xf32>
%318 = tensor.empty() : tensor<4x8xf32>
%319 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%312 : tensor<4x8xf32>) outs(%318 : tensor<4x8xf32>) {
^bb0(%x: f32, %_: f32):
%317 = math.sqrt %x : f32
linalg.yield %317 : f32
} -> tensor<4x8xf32>
%321 = tensor.empty() : tensor<1x1xf32>
%322 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%316 : tensor<1x1xf32>) outs(%321 : tensor<1x1xf32>) {
^bb0(%x: f32, %_: f32):
%320 = math.sqrt %x : f32
linalg.yield %320 : f32
} -> tensor<1x1xf32>
%323 = tensor.empty() : tensor<4x8xf32>
%324 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%322 : tensor<1x1xf32>) outs(%323 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%325 = tensor.empty() : tensor<4x8xf32>
%326 = linalg.div ins(%319, %324 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%325 : tensor<4x8xf32>) -> tensor<4x8xf32>
%327 = tensor.empty() : tensor<4x8xf32>
%328 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%282 : tensor<1x1xf32>) outs(%327 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%329 = tensor.empty() : tensor<4x8xf32>
%330 = linalg.add ins(%326, %328 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%329 : tensor<4x8xf32>) -> tensor<4x8xf32>
%331 = tensor.empty() : tensor<1x1xf32>
%332 = linalg.div ins(%279, %314 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%331 : tensor<1x1xf32>) -> tensor<1x1xf32>
%333 = tensor.empty() : tensor<4x8xf32>
%334 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%332 : tensor<1x1xf32>) outs(%333 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%335 = tensor.empty() : tensor<4x8xf32>
%336 = linalg.div ins(%304, %330 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%335 : tensor<4x8xf32>) -> tensor<4x8xf32>
%337 = tensor.empty() : tensor<4x8xf32>
%338 = linalg.mul ins(%334, %336 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%337 : tensor<4x8xf32>) -> tensor<4x8xf32>
%339 = tensor.empty() : tensor<1x1xf32>
%340 = linalg.mul ins(%279, %283 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%339 : tensor<1x1xf32>) -> tensor<1x1xf32>
%341 = tensor.empty() : tensor<1x1xf32>
%342 = linalg.sub ins(%286, %340 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%341 : tensor<1x1xf32>) -> tensor<1x1xf32>
%343 = tensor.empty() : tensor<4x8xf32>
%344 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%342 : tensor<1x1xf32>) outs(%343 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%345 = tensor.empty() : tensor<4x8xf32>
%346 = linalg.mul ins(%275, %344 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%345 : tensor<4x8xf32>) -> tensor<4x8xf32>
%347 = tensor.empty() : tensor<4x8xf32>
%348 = linalg.sub ins(%346, %338 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%347 : tensor<4x8xf32>) -> tensor<4x8xf32>
return %348, %304, %312 : tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>
}
func.func @fn_val_349(%350: tensor<1x4xi64>, %351: tensor<1x4xi64>, %352: tensor<8x4xf32>, %353: tensor<1x1x4x4xf32>, %354: tensor<1x1x1x1xf32>, %355: tensor<4x4xf32>, %356: tensor<4x4xf32>, %357: tensor<4x4xf32>, %358: tensor<4x4xf32>, %359: tensor<4x8xf32>) -> tensor<1x4xf32> attributes {llvm.emit_c_interface} {
%360 = arith.constant 0 : i64
%361 = arith.constant 8 : i64
%362 = tensor.empty() : tensor<1x4xi64>
%363 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%350 : tensor<1x4xi64>) outs(%362 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%364 = arith.cmpi sge, %candidate, %360 : i64
%365 = arith.cmpi slt, %candidate, %361 : i64
%366 = arith.andi %364, %365 : i1
cf.assert %366, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%367 = tensor.empty() : tensor<1x4x4xf32>
%368 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%363 : tensor<1x4xi64>) outs(%367 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%369 = arith.index_cast %idx : i64 to index
%370 = linalg.index 2 : index
%371 = tensor.extract %352[%369, %370] : tensor<8x4xf32>
linalg.yield %371 : f32
} -> tensor<1x4x4xf32>
%372 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%373 = tensor.reshape %355(%372) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%374 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%375 = tensor.reshape %356(%374) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%376 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%377 = tensor.reshape %357(%376) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%378 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%379 = tensor.reshape %358(%378) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%380 = arith.constant dense<[1, 4, 8]> : tensor<3xi32>
%381 = tensor.reshape %359(%380) : (tensor<4x8xf32>, tensor<3xi32>) -> tensor<1x4x8xf32>
%382 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%383 = linalg.batch_matmul ins(%368, %373 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%382 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%384 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%385 = linalg.batch_matmul ins(%368, %375 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%384 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%386 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%387 = linalg.batch_matmul ins(%368, %377 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%386 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%388 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%389 = tensor.reshape %383(%388) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%390 = tensor.empty() : tensor<1x2x4x2xf32>
%391 = linalg.transpose ins(%389 : tensor<1x4x2x2xf32>) outs(%390 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%392 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%393 = tensor.reshape %385(%392) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%394 = tensor.empty() : tensor<1x2x2x4xf32>
%395 = linalg.transpose ins(%393 : tensor<1x4x2x2xf32>) outs(%394 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%396 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%397 = tensor.reshape %387(%396) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%398 = tensor.empty() : tensor<1x2x4x2xf32>
%399 = linalg.transpose ins(%397 : tensor<1x4x2x2xf32>) outs(%398 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%400 = tensor.collapse_shape %391 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%401 = tensor.collapse_shape %395 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%402 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%403 = linalg.batch_matmul ins(%400, %401 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%402 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%404 = tensor.expand_shape %403 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%405 = tensor.empty() : tensor<1x2x4x4xf32>
%406 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, 0, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%354 : tensor<1x1x1x1xf32>) outs(%405 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%407 = tensor.empty() : tensor<1x2x4x4xf32>
%408 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%353 : tensor<1x1x4x4xf32>) outs(%407 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%409 = tensor.empty() : tensor<1x2x4x4xf32>
%410 = linalg.mul ins(%404, %406 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%409 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%411 = tensor.empty() : tensor<1x2x4x4xf32>
%412 = linalg.add ins(%410, %408 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%411 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%415 = arith.constant 0xFF800000 : f32
%413 = tensor.empty() : tensor<1x2x4x1xf32>
%414 = linalg.fill ins(%415 : f32) outs(%413 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%416 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%412 : tensor<1x2x4x4xf32>) outs(%414 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%417 = arith.maximumf %a0, %a1 : f32
linalg.yield %417 : f32
} -> tensor<1x2x4x1xf32>
%418 = tensor.empty() : tensor<1x2x4x4xf32>
%419 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%416 : tensor<1x2x4x1xf32>) outs(%418 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%420 = tensor.empty() : tensor<1x2x4x4xf32>
%421 = linalg.sub ins(%412, %419 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%420 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%422 = tensor.empty() : tensor<1x2x4x4xf32>
%423 = linalg.exp ins(%421 : tensor<1x2x4x4xf32>) outs(%422 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%426 = arith.constant 0xFF800000 : f32
%424 = tensor.empty() : tensor<1x2x4x1xf32>
%425 = linalg.fill ins(%426 : f32) outs(%424 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%427 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%412 : tensor<1x2x4x4xf32>) outs(%425 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%428 = arith.maximumf %a0, %a1 : f32
linalg.yield %428 : f32
} -> tensor<1x2x4x1xf32>
%429 = tensor.empty() : tensor<1x2x4x4xf32>
%430 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%427 : tensor<1x2x4x1xf32>) outs(%429 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%431 = tensor.empty() : tensor<1x2x4x4xf32>
%432 = linalg.sub ins(%412, %430 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%431 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%433 = tensor.empty() : tensor<1x2x4x4xf32>
%434 = linalg.exp ins(%432 : tensor<1x2x4x4xf32>) outs(%433 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%437 = arith.constant 0.000000e+00 : f32
%435 = tensor.empty() : tensor<1x2x4x1xf32>
%436 = linalg.fill ins(%437 : f32) outs(%435 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%438 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%434 : tensor<1x2x4x4xf32>) outs(%436 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%439 = arith.addf %a0, %a1 : f32
linalg.yield %439 : f32
} -> tensor<1x2x4x1xf32>
%440 = tensor.empty() : tensor<1x2x4x4xf32>
%441 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%438 : tensor<1x2x4x1xf32>) outs(%440 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%442 = tensor.empty() : tensor<1x2x4x4xf32>
%443 = linalg.div ins(%423, %441 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%442 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%444 = tensor.collapse_shape %443 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%445 = tensor.collapse_shape %399 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%446 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%447 = linalg.batch_matmul ins(%444, %445 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%446 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%448 = tensor.expand_shape %447 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%449 = tensor.empty() : tensor<1x4x2x2xf32>
%450 = linalg.transpose ins(%448 : tensor<1x2x4x2xf32>) outs(%449 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%451 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%452 = tensor.reshape %450(%451) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%453 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%454 = linalg.batch_matmul ins(%452, %379 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%453 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%455 = tensor.empty() : tensor<1x4x4xf32>
%456 = linalg.add ins(%368, %454 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%455 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%457 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%458 = linalg.batch_matmul ins(%456, %381 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%457 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%461 = arith.constant 0xFF800000 : f32
%459 = tensor.empty() : tensor<1x4x1xf32>
%460 = linalg.fill ins(%461 : f32) outs(%459 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%462 = 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(%458 : tensor<1x4x8xf32>) outs(%460 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%463 = arith.maximumf %a0, %a1 : f32
linalg.yield %463 : f32
} -> tensor<1x4x1xf32>
%464 = tensor.empty() : tensor<1x4x8xf32>
%465 = 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(%462 : tensor<1x4x1xf32>) outs(%464 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%466 = tensor.empty() : tensor<1x4x8xf32>
%467 = linalg.sub ins(%458, %465 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%466 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%470 = arith.constant 0xFF800000 : f32
%468 = tensor.empty() : tensor<1x4x1xf32>
%469 = linalg.fill ins(%470 : f32) outs(%468 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%471 = 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(%458 : tensor<1x4x8xf32>) outs(%469 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%472 = arith.maximumf %a0, %a1 : f32
linalg.yield %472 : f32
} -> tensor<1x4x1xf32>
%473 = tensor.empty() : tensor<1x4x8xf32>
%474 = 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(%471 : tensor<1x4x1xf32>) outs(%473 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%475 = tensor.empty() : tensor<1x4x8xf32>
%476 = linalg.sub ins(%458, %474 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%475 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%477 = tensor.empty() : tensor<1x4x8xf32>
%478 = linalg.exp ins(%476 : tensor<1x4x8xf32>) outs(%477 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%481 = arith.constant 0.000000e+00 : f32
%479 = tensor.empty() : tensor<1x4x1xf32>
%480 = linalg.fill ins(%481 : f32) outs(%479 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%482 = 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(%478 : tensor<1x4x8xf32>) outs(%480 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%483 = arith.addf %a0, %a1 : f32
linalg.yield %483 : f32
} -> tensor<1x4x1xf32>
%484 = tensor.empty() : tensor<1x4x8xf32>
%485 = 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(%482 : tensor<1x4x1xf32>) outs(%484 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%486 = tensor.empty() : tensor<1x4x8xf32>
%487 = linalg.log ins(%485 : tensor<1x4x8xf32>) outs(%486 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%488 = tensor.empty() : tensor<1x4x8xf32>
%489 = linalg.sub ins(%467, %487 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%488 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%490 = arith.constant 0 : i64
%491 = arith.constant 8 : i64
%492 = tensor.empty() : tensor<1x4xi64>
%493 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%351 : tensor<1x4xi64>) outs(%492 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%494 = arith.cmpi sge, %candidate, %490 : i64
%495 = arith.cmpi slt, %candidate, %491 : i64
%496 = arith.andi %494, %495 : i1
cf.assert %496, "take_along_axis index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%497 = tensor.empty() : tensor<1x4xf32>
%498 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%493 : tensor<1x4xi64>) outs(%497 : tensor<1x4xf32>) {
^bb0(%idx: i64, %_: f32):
%499 = linalg.index 0 : index
%500 = linalg.index 1 : index
%501 = arith.index_cast %idx : i64 to index
%502 = tensor.extract %489[%499, %500, %501] : tensor<1x4x8xf32>
linalg.yield %502 : f32
} -> tensor<1x4xf32>
%503 = tensor.empty() : tensor<1x4xf32>
%504 = linalg.negf ins(%498 : tensor<1x4xf32>) outs(%503 : tensor<1x4xf32>) -> tensor<1x4xf32>
return %504 : tensor<1x4xf32>
}
func.func @fn_val_505(%506: tensor<1x4xi64>, %507: tensor<1x4xi64>, %508: tensor<8x4xf32>, %509: tensor<1x1x4x4xf32>, %510: tensor<1x1x1x1xf32>, %511: tensor<4x4xf32>, %512: tensor<4x4xf32>, %513: tensor<4x4xf32>, %514: tensor<4x4xf32>, %515: tensor<4x8xf32>, %516: tensor<8x4xf32>, %517: tensor<8x4xf32>, %518: tensor<4x4xf32>, %519: tensor<4x4xf32>, %520: tensor<4x4xf32>, %521: tensor<4x4xf32>, %522: tensor<4x4xf32>, %523: tensor<4x4xf32>, %524: tensor<4x4xf32>, %525: tensor<4x4xf32>, %526: tensor<4x8xf32>, %527: tensor<4x8xf32>, %528: tensor<1x1xf32>, %529: tensor<1x1xf32>, %530: tensor<1x1xf32>, %531: tensor<1x1xf32>, %532: tensor<1x1xf32>, %533: tensor<1x1xf32>, %534: tensor<1x1xf32>, %535: tensor<1x1xf32>) -> (tensor<1x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>) attributes {llvm.emit_c_interface} {
%536 = arith.constant 0 : i64
%537 = arith.constant 8 : i64
%538 = tensor.empty() : tensor<1x4xi64>
%539 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%506 : tensor<1x4xi64>) outs(%538 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%540 = arith.cmpi sge, %candidate, %536 : i64
%541 = arith.cmpi slt, %candidate, %537 : i64
%542 = arith.andi %540, %541 : i1
cf.assert %542, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%543 = tensor.empty() : tensor<1x4x4xf32>
%544 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%539 : tensor<1x4xi64>) outs(%543 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%545 = arith.index_cast %idx : i64 to index
%546 = linalg.index 2 : index
%547 = tensor.extract %508[%545, %546] : tensor<8x4xf32>
linalg.yield %547 : f32
} -> tensor<1x4x4xf32>
%548 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%549 = tensor.reshape %511(%548) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%550 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%551 = tensor.reshape %512(%550) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%552 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%553 = tensor.reshape %513(%552) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%554 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%555 = tensor.reshape %514(%554) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%556 = arith.constant dense<[1, 4, 8]> : tensor<3xi32>
%557 = tensor.reshape %515(%556) : (tensor<4x8xf32>, tensor<3xi32>) -> tensor<1x4x8xf32>
%558 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%559 = linalg.batch_matmul ins(%544, %549 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%558 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%560 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%561 = linalg.batch_matmul ins(%544, %551 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%560 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%562 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%563 = linalg.batch_matmul ins(%544, %553 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%562 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%564 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%565 = tensor.reshape %559(%564) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%566 = tensor.empty() : tensor<1x2x4x2xf32>
%567 = linalg.transpose ins(%565 : tensor<1x4x2x2xf32>) outs(%566 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%568 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%569 = tensor.reshape %561(%568) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%570 = tensor.empty() : tensor<1x2x2x4xf32>
%571 = linalg.transpose ins(%569 : tensor<1x4x2x2xf32>) outs(%570 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%572 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%573 = tensor.reshape %563(%572) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%574 = tensor.empty() : tensor<1x2x4x2xf32>
%575 = linalg.transpose ins(%573 : tensor<1x4x2x2xf32>) outs(%574 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%576 = tensor.collapse_shape %567 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%577 = tensor.collapse_shape %571 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%578 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%579 = linalg.batch_matmul ins(%576, %577 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%578 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%580 = tensor.expand_shape %579 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%581 = tensor.empty() : tensor<1x2x4x4xf32>
%582 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, 0, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%510 : tensor<1x1x1x1xf32>) outs(%581 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%583 = tensor.empty() : tensor<1x2x4x4xf32>
%584 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%509 : tensor<1x1x4x4xf32>) outs(%583 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%585 = tensor.empty() : tensor<1x2x4x4xf32>
%586 = linalg.mul ins(%580, %582 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%585 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%587 = tensor.empty() : tensor<1x2x4x4xf32>
%588 = linalg.add ins(%586, %584 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%587 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%591 = arith.constant 0xFF800000 : f32
%589 = tensor.empty() : tensor<1x2x4x1xf32>
%590 = linalg.fill ins(%591 : f32) outs(%589 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%592 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%588 : tensor<1x2x4x4xf32>) outs(%590 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%593 = arith.maximumf %a0, %a1 : f32
linalg.yield %593 : f32
} -> tensor<1x2x4x1xf32>
%594 = tensor.empty() : tensor<1x2x4x4xf32>
%595 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%592 : tensor<1x2x4x1xf32>) outs(%594 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%596 = tensor.empty() : tensor<1x2x4x4xf32>
%597 = linalg.sub ins(%588, %595 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%596 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%598 = tensor.empty() : tensor<1x2x4x4xf32>
%599 = linalg.exp ins(%597 : tensor<1x2x4x4xf32>) outs(%598 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%602 = arith.constant 0xFF800000 : f32
%600 = tensor.empty() : tensor<1x2x4x1xf32>
%601 = linalg.fill ins(%602 : f32) outs(%600 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%603 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%588 : tensor<1x2x4x4xf32>) outs(%601 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%604 = arith.maximumf %a0, %a1 : f32
linalg.yield %604 : f32
} -> tensor<1x2x4x1xf32>
%605 = tensor.empty() : tensor<1x2x4x4xf32>
%606 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%603 : tensor<1x2x4x1xf32>) outs(%605 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%607 = tensor.empty() : tensor<1x2x4x4xf32>
%608 = linalg.sub ins(%588, %606 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%607 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%609 = tensor.empty() : tensor<1x2x4x4xf32>
%610 = linalg.exp ins(%608 : tensor<1x2x4x4xf32>) outs(%609 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%613 = arith.constant 0.000000e+00 : f32
%611 = tensor.empty() : tensor<1x2x4x1xf32>
%612 = linalg.fill ins(%613 : f32) outs(%611 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%614 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%610 : tensor<1x2x4x4xf32>) outs(%612 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%615 = arith.addf %a0, %a1 : f32
linalg.yield %615 : f32
} -> tensor<1x2x4x1xf32>
%616 = tensor.empty() : tensor<1x2x4x4xf32>
%617 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%614 : tensor<1x2x4x1xf32>) outs(%616 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%618 = tensor.empty() : tensor<1x2x4x4xf32>
%619 = linalg.div ins(%599, %617 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%618 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%620 = tensor.collapse_shape %619 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%621 = tensor.collapse_shape %575 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%622 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%623 = linalg.batch_matmul ins(%620, %621 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%622 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%624 = tensor.expand_shape %623 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%625 = tensor.empty() : tensor<1x4x2x2xf32>
%626 = linalg.transpose ins(%624 : tensor<1x2x4x2xf32>) outs(%625 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%627 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%628 = tensor.reshape %626(%627) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%629 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%630 = linalg.batch_matmul ins(%628, %555 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%629 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%631 = tensor.empty() : tensor<1x4x4xf32>
%632 = linalg.add ins(%544, %630 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%631 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%633 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%634 = linalg.batch_matmul ins(%632, %557 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%633 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%637 = arith.constant 0xFF800000 : f32
%635 = tensor.empty() : tensor<1x4x1xf32>
%636 = linalg.fill ins(%637 : f32) outs(%635 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%638 = 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(%634 : tensor<1x4x8xf32>) outs(%636 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%639 = arith.maximumf %a0, %a1 : f32
linalg.yield %639 : f32
} -> tensor<1x4x1xf32>
%640 = tensor.empty() : tensor<1x4x8xf32>
%641 = 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(%638 : tensor<1x4x1xf32>) outs(%640 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%642 = tensor.empty() : tensor<1x4x8xf32>
%643 = linalg.sub ins(%634, %641 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%642 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%646 = arith.constant 0xFF800000 : f32
%644 = tensor.empty() : tensor<1x4x1xf32>
%645 = linalg.fill ins(%646 : f32) outs(%644 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%647 = 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(%634 : tensor<1x4x8xf32>) outs(%645 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%648 = arith.maximumf %a0, %a1 : f32
linalg.yield %648 : f32
} -> tensor<1x4x1xf32>
%649 = tensor.empty() : tensor<1x4x8xf32>
%650 = 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(%647 : tensor<1x4x1xf32>) outs(%649 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%651 = tensor.empty() : tensor<1x4x8xf32>
%652 = linalg.sub ins(%634, %650 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%651 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%653 = tensor.empty() : tensor<1x4x8xf32>
%654 = linalg.exp ins(%652 : tensor<1x4x8xf32>) outs(%653 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%657 = arith.constant 0.000000e+00 : f32
%655 = tensor.empty() : tensor<1x4x1xf32>
%656 = linalg.fill ins(%657 : f32) outs(%655 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%658 = 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(%654 : tensor<1x4x8xf32>) outs(%656 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%659 = arith.addf %a0, %a1 : f32
linalg.yield %659 : f32
} -> tensor<1x4x1xf32>
%660 = tensor.empty() : tensor<1x4x8xf32>
%661 = 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(%658 : tensor<1x4x1xf32>) outs(%660 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%662 = tensor.empty() : tensor<1x4x8xf32>
%663 = linalg.log ins(%661 : tensor<1x4x8xf32>) outs(%662 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%664 = tensor.empty() : tensor<1x4x8xf32>
%665 = linalg.sub ins(%643, %663 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%664 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%666 = arith.constant 0 : i64
%667 = arith.constant 8 : i64
%668 = tensor.empty() : tensor<1x4xi64>
%669 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%507 : tensor<1x4xi64>) outs(%668 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%670 = arith.cmpi sge, %candidate, %666 : i64
%671 = arith.cmpi slt, %candidate, %667 : i64
%672 = arith.andi %670, %671 : i1
cf.assert %672, "take_along_axis index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%673 = tensor.empty() : tensor<1x4xf32>
%674 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%669 : tensor<1x4xi64>) outs(%673 : tensor<1x4xf32>) {
^bb0(%idx: i64, %_: f32):
%675 = linalg.index 0 : index
%676 = linalg.index 1 : index
%677 = arith.index_cast %idx : i64 to index
%678 = tensor.extract %665[%675, %676, %677] : tensor<1x4x8xf32>
linalg.yield %678 : f32
} -> tensor<1x4xf32>
%679 = tensor.empty() : tensor<1x4xf32>
%680 = linalg.negf ins(%674 : tensor<1x4xf32>) outs(%679 : tensor<1x4xf32>) -> tensor<1x4xf32>
%681 = arith.constant dense<1.000000e+00> : tensor<1x4xf32>
%682 = arith.constant dense<1.000000e+00> : tensor<1x4xf32>
%683 = tensor.empty() : tensor<1x4xf32>
%684 = linalg.negf ins(%682 : tensor<1x4xf32>) outs(%683 : tensor<1x4xf32>) -> tensor<1x4xf32>
%685 = arith.constant 0 : i64
%686 = arith.constant 8 : i64
%687 = tensor.empty() : tensor<1x4xi64>
%688 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%507 : tensor<1x4xi64>) outs(%687 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%689 = arith.cmpi sge, %candidate, %685 : i64
%690 = arith.cmpi slt, %candidate, %686 : i64
%691 = arith.andi %689, %690 : i1
cf.assert %691, "scatter_along_axis index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%692 = tensor.empty() : tensor<1x4x8xf32>
%693 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%688, %684 : tensor<1x4xi64>, tensor<1x4xf32>) outs(%692 : tensor<1x4x8xf32>) {
^bb0(%idx: i64, %upd: f32, %_: f32):
%694 = linalg.index 2 : index
%695 = arith.index_cast %694 : index to i64
%696 = arith.cmpi eq, %idx, %695 : i64
%697 = arith.constant 0.000000e+00 : f32
%698 = arith.select %696, %upd, %697 : f32
linalg.yield %698 : f32
} -> tensor<1x4x8xf32>
%699 = tensor.empty() : tensor<1x4x8xf32>
%700 = linalg.exp ins(%665 : tensor<1x4x8xf32>) outs(%699 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%703 = arith.constant 0.000000e+00 : f32
%701 = tensor.empty() : tensor<1x4x1xf32>
%702 = linalg.fill ins(%703 : f32) outs(%701 : tensor<1x4x1xf32>) -> tensor<1x4x1xf32>
%704 = 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(%693 : tensor<1x4x8xf32>) outs(%702 : tensor<1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%705 = arith.addf %a0, %a1 : f32
linalg.yield %705 : f32
} -> tensor<1x4x1xf32>
%706 = tensor.empty() : tensor<1x4x8xf32>
%707 = 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(%704 : tensor<1x4x1xf32>) outs(%706 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%708 = tensor.empty() : tensor<1x4x8xf32>
%709 = linalg.mul ins(%707, %700 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%708 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%710 = tensor.empty() : tensor<1x4x8xf32>
%711 = linalg.sub ins(%693, %709 : tensor<1x4x8xf32>, tensor<1x4x8xf32>) outs(%710 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%712 = tensor.empty() : tensor<1x4x4xf32>
%713 = linalg.transpose ins(%632 : tensor<1x4x4xf32>) outs(%712 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%714 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%715 = linalg.batch_matmul ins(%713, %711 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%714 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%716 = tensor.empty() : tensor<1x8x4xf32>
%717 = linalg.transpose ins(%557 : tensor<1x4x8xf32>) outs(%716 : tensor<1x8x4xf32>) permutation = [0, 2, 1]
%718 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%719 = linalg.batch_matmul ins(%711, %717 : tensor<1x4x8xf32>, tensor<1x8x4xf32>) outs(%718 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%720 = tensor.empty() : tensor<1x4x4xf32>
%721 = linalg.transpose ins(%628 : tensor<1x4x4xf32>) outs(%720 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%722 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%723 = linalg.batch_matmul ins(%721, %719 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%722 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%724 = tensor.empty() : tensor<1x4x4xf32>
%725 = linalg.transpose ins(%555 : tensor<1x4x4xf32>) outs(%724 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%726 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%727 = linalg.batch_matmul ins(%719, %725 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%726 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%728 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%729 = tensor.reshape %727(%728) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%730 = tensor.empty() : tensor<1x2x4x2xf32>
%731 = linalg.transpose ins(%729 : tensor<1x4x2x2xf32>) outs(%730 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%732 = tensor.empty() : tensor<1x2x4x4xf32>
%733 = linalg.transpose ins(%619 : tensor<1x2x4x4xf32>) outs(%732 : tensor<1x2x4x4xf32>) permutation = [0, 1, 3, 2]
%734 = tensor.collapse_shape %733 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%735 = tensor.collapse_shape %731 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%736 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%737 = linalg.batch_matmul ins(%734, %735 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%736 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%738 = tensor.expand_shape %737 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%739 = tensor.empty() : tensor<1x2x2x4xf32>
%740 = linalg.transpose ins(%575 : tensor<1x2x4x2xf32>) outs(%739 : tensor<1x2x2x4xf32>) permutation = [0, 1, 3, 2]
%741 = tensor.collapse_shape %731 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%742 = tensor.collapse_shape %740 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%743 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%744 = linalg.batch_matmul ins(%741, %742 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%743 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%745 = tensor.expand_shape %744 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%746 = tensor.empty() : tensor<1x2x4x4xf32>
%747 = linalg.mul ins(%619, %745 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%746 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%750 = arith.constant 0.000000e+00 : f32
%748 = tensor.empty() : tensor<1x2x4x1xf32>
%749 = linalg.fill ins(%750 : f32) outs(%748 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%751 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%747 : tensor<1x2x4x4xf32>) outs(%749 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%752 = arith.addf %a0, %a1 : f32
linalg.yield %752 : f32
} -> tensor<1x2x4x1xf32>
%753 = tensor.empty() : tensor<1x2x4x4xf32>
%754 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, d1, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%751 : tensor<1x2x4x1xf32>) outs(%753 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%755 = tensor.empty() : tensor<1x2x4x4xf32>
%756 = linalg.sub ins(%745, %754 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%755 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%757 = tensor.empty() : tensor<1x2x4x4xf32>
%758 = linalg.mul ins(%619, %756 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%757 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%759 = tensor.empty() : tensor<1x2x4x4xf32>
%760 = linalg.mul ins(%758, %582 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%759 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%761 = tensor.empty() : tensor<1x2x4x4xf32>
%762 = linalg.mul ins(%758, %580 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%761 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%765 = arith.constant 0.000000e+00 : f32
%763 = tensor.empty() : tensor<1x1x4x4xf32>
%764 = linalg.fill ins(%765 : f32) outs(%763 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%766 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, 0, d2, d3)>],
iterator_types = ["parallel", "reduction", "parallel", "parallel"]
} ins(%758 : tensor<1x2x4x4xf32>) outs(%764 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
%767 = arith.addf %a0, %a1 : f32
linalg.yield %767 : f32
} -> tensor<1x1x4x4xf32>
%770 = arith.constant 0.000000e+00 : f32
%768 = tensor.empty() : tensor<1x2x4x1xf32>
%769 = linalg.fill ins(%770 : f32) outs(%768 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%771 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
} ins(%762 : tensor<1x2x4x4xf32>) outs(%769 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%772 = arith.addf %a0, %a1 : f32
linalg.yield %772 : f32
} -> tensor<1x2x4x1xf32>
%775 = arith.constant 0.000000e+00 : f32
%773 = tensor.empty() : tensor<1x2x1x1xf32>
%774 = linalg.fill ins(%775 : f32) outs(%773 : tensor<1x2x1x1xf32>) -> tensor<1x2x1x1xf32>
%776 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, 0, d3)>],
iterator_types = ["parallel", "parallel", "reduction", "parallel"]
} ins(%771 : tensor<1x2x4x1xf32>) outs(%774 : tensor<1x2x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%777 = arith.addf %a0, %a1 : f32
linalg.yield %777 : f32
} -> tensor<1x2x1x1xf32>
%780 = arith.constant 0.000000e+00 : f32
%778 = tensor.empty() : tensor<1x1x1x1xf32>
%779 = linalg.fill ins(%780 : f32) outs(%778 : tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
%781 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, 0, d2, d3)>],
iterator_types = ["parallel", "reduction", "parallel", "parallel"]
} ins(%776 : tensor<1x2x1x1xf32>) outs(%779 : tensor<1x1x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%782 = arith.addf %a0, %a1 : f32
linalg.yield %782 : f32
} -> tensor<1x1x1x1xf32>
%783 = tensor.empty() : tensor<1x2x2x4xf32>
%784 = linalg.transpose ins(%567 : tensor<1x2x4x2xf32>) outs(%783 : tensor<1x2x2x4xf32>) permutation = [0, 1, 3, 2]
%785 = tensor.collapse_shape %784 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%786 = tensor.collapse_shape %760 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%787 = arith.constant dense<0.000000e+00> : tensor<2x2x4xf32>
%788 = linalg.batch_matmul ins(%785, %786 : tensor<2x2x4xf32>, tensor<2x4x4xf32>) outs(%787 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
%789 = tensor.expand_shape %788 [[0, 1], [2], [3]] : tensor<2x2x4xf32> into tensor<1x2x2x4xf32>
%790 = tensor.empty() : tensor<1x2x4x2xf32>
%791 = linalg.transpose ins(%571 : tensor<1x2x2x4xf32>) outs(%790 : tensor<1x2x4x2xf32>) permutation = [0, 1, 3, 2]
%792 = tensor.collapse_shape %760 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%793 = tensor.collapse_shape %791 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%794 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%795 = linalg.batch_matmul ins(%792, %793 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%794 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%796 = tensor.expand_shape %795 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%797 = tensor.empty() : tensor<1x4x2x2xf32>
%798 = linalg.transpose ins(%738 : tensor<1x2x4x2xf32>) outs(%797 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%799 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%800 = tensor.reshape %798(%799) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%801 = tensor.empty() : tensor<1x4x2x2xf32>
%802 = linalg.transpose ins(%789 : tensor<1x2x2x4xf32>) outs(%801 : tensor<1x4x2x2xf32>) permutation = [0, 3, 1, 2]
%803 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%804 = tensor.reshape %802(%803) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%805 = tensor.empty() : tensor<1x4x2x2xf32>
%806 = linalg.transpose ins(%796 : tensor<1x2x4x2xf32>) outs(%805 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%807 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%808 = tensor.reshape %806(%807) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%809 = tensor.empty() : tensor<1x4x4xf32>
%810 = linalg.transpose ins(%544 : tensor<1x4x4xf32>) outs(%809 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%811 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%812 = linalg.batch_matmul ins(%810, %800 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%811 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%813 = tensor.empty() : tensor<1x4x4xf32>
%814 = linalg.transpose ins(%553 : tensor<1x4x4xf32>) outs(%813 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%815 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%816 = linalg.batch_matmul ins(%800, %814 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%815 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%817 = tensor.empty() : tensor<1x4x4xf32>
%818 = linalg.transpose ins(%544 : tensor<1x4x4xf32>) outs(%817 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%819 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%820 = linalg.batch_matmul ins(%818, %804 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%819 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%821 = tensor.empty() : tensor<1x4x4xf32>
%822 = linalg.transpose ins(%551 : tensor<1x4x4xf32>) outs(%821 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%823 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%824 = linalg.batch_matmul ins(%804, %822 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%823 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%825 = tensor.empty() : tensor<1x4x4xf32>
%826 = linalg.transpose ins(%544 : tensor<1x4x4xf32>) outs(%825 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%827 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%828 = linalg.batch_matmul ins(%826, %808 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%827 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%829 = tensor.empty() : tensor<1x4x4xf32>
%830 = linalg.transpose ins(%549 : tensor<1x4x4xf32>) outs(%829 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%831 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%832 = linalg.batch_matmul ins(%808, %830 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%831 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%833 = arith.constant dense<[4, 8]> : tensor<2xi32>
%834 = tensor.reshape %715(%833) : (tensor<1x4x8xf32>, tensor<2xi32>) -> tensor<4x8xf32>
%835 = arith.constant dense<[4, 4]> : tensor<2xi32>
%836 = tensor.reshape %723(%835) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%837 = arith.constant dense<[4, 4]> : tensor<2xi32>
%838 = tensor.reshape %812(%837) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%839 = arith.constant dense<[4, 4]> : tensor<2xi32>
%840 = tensor.reshape %820(%839) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%841 = arith.constant dense<[4, 4]> : tensor<2xi32>
%842 = tensor.reshape %828(%841) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%843 = tensor.empty() : tensor<1x4x4xf32>
%844 = linalg.add ins(%719, %816 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%843 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%845 = tensor.empty() : tensor<1x4x4xf32>
%846 = linalg.add ins(%844, %824 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%845 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%847 = tensor.empty() : tensor<1x4x4xf32>
%848 = linalg.add ins(%846, %832 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%847 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%849 = arith.constant 0 : i64
%850 = arith.constant 8 : i64
%851 = tensor.empty() : tensor<1x4xi64>
%852 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%506 : tensor<1x4xi64>) outs(%851 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%853 = arith.cmpi sge, %candidate, %849 : i64
%854 = arith.cmpi slt, %candidate, %850 : i64
%855 = arith.andi %853, %854 : i1
cf.assert %855, "scatter_add index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%856 = tensor.collapse_shape %852 [[0, 1]] : tensor<1x4xi64> into tensor<4xi64>
%857 = tensor.collapse_shape %848 [[0, 1], [2]] : tensor<1x4x4xf32> into tensor<4x4xf32>
%858 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%859 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%856, %857 : tensor<4xi64>, tensor<4x4xf32>) outs(%858 : tensor<8x4xf32>) {
^bb0(%idx: i64, %upd: f32, %acc: f32):
%860 = linalg.index 0 : index
%861 = arith.index_cast %860 : index to i64
%862 = arith.cmpi eq, %idx, %861 : i64
%863 = arith.constant 0.000000e+00 : f32
%864 = arith.select %862, %upd, %863 : f32
%865 = arith.addf %acc, %864 : f32
linalg.yield %865 : f32
} -> tensor<8x4xf32>
%866, %867, %868 = func.call @fn_val_124(%508, %859, %516, %517, %528, %529, %530, %531, %532, %533, %534, %535) : (tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>)
%869, %870, %871 = func.call @fn_val_199(%511, %842, %518, %519, %528, %529, %530, %531, %532, %533, %534, %535) : (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>)
%872, %873, %874 = func.call @fn_val_199(%512, %840, %520, %521, %528, %529, %530, %531, %532, %533, %534, %535) : (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>)
%875, %876, %877 = func.call @fn_val_199(%513, %838, %522, %523, %528, %529, %530, %531, %532, %533, %534, %535) : (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>)
%878, %879, %880 = func.call @fn_val_199(%514, %836, %524, %525, %528, %529, %530, %531, %532, %533, %534, %535) : (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>)
%881, %882, %883 = func.call @fn_val_274(%515, %834, %526, %527, %528, %529, %530, %531, %532, %533, %534, %535) : (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>)
return %680, %866, %867, %868, %869, %870, %871, %872, %873, %874, %875, %876, %877, %878, %879, %880, %881, %882, %883 : tensor<1x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>
}
func.func @main() -> (tensor<1x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>) {
%884 = arith.constant dense<0> : tensor<1x4xi64>
%885 = arith.constant dense<0> : tensor<1x4xi64>
%886 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%887 = arith.constant dense<0.000000e+00> : tensor<1x1x4x4xf32>
%888 = arith.constant dense<0.000000e+00> : tensor<1x1x1x1xf32>
%889 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%890 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%891 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%892 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%893 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%894 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%895 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%896 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%897 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%898 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%899 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%900 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%901 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%902 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%903 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%904 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%905 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%906 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%907 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%908 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%909 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%910 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%911 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%912 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%913 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%914, %915, %916, %917, %918, %919, %920, %921, %922, %923, %924, %925, %926, %927, %928, %929, %930, %931, %932 = func.call @fn_val_505(%884, %885, %886, %887, %888, %889, %890, %891, %892, %893, %894, %895, %896, %897, %898, %899, %900, %901, %902, %903, %904, %905, %906, %907, %908, %909, %910, %911, %912, %913) : (tensor<1x4xi64>, tensor<1x4xi64>, tensor<8x4xf32>, tensor<1x1x4x4xf32>, tensor<1x1x1x1xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<1x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>)
return %914, %915, %916, %917, %918, %919, %920, %921, %922, %923, %924, %925, %926, %927, %928, %929, %930, %931, %932 : tensor<1x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<8x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>
}
}