Example
storylm0_ad_step.kina
let storylm0_fwd = fn(
token_ids: 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, 8], @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
logits
} in
let step = fn(
token_ids: 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>,
lr: tensor<F32, [1, 1], @host>
) -> tensor<F32, [8, 4], @host> {
let (logits, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm) =
handle storylm0_fwd(token_ids, token_embedding, causal_mask, attention_scale, wq, wk, wv, wo, lm_head) with reverseMode in
let lr_b = lr.broadcast([8, 4]) in
let update = lr_b * g_emb in
token_embedding - update
} in
let token_ids = 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 lr = tensor<F32, [1, 1], @host> in
step(token_ids, token_embedding, causal_mask, attention_scale, wq, wk, wv, wo, lm_head, lr)
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<1x4xi64>, %126: tensor<8x4xf32>, %127: tensor<1x1x4x4xf32>, %128: tensor<1x1x1x1xf32>, %129: tensor<4x4xf32>, %130: tensor<4x4xf32>, %131: tensor<4x4xf32>, %132: tensor<4x4xf32>, %133: tensor<4x8xf32>) -> tensor<1x4x8xf32> attributes {llvm.emit_c_interface} {
%134 = arith.constant 0 : i64
%135 = arith.constant 8 : i64
%136 = tensor.empty() : tensor<1x4xi64>
%137 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%125 : tensor<1x4xi64>) outs(%136 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%138 = arith.cmpi sge, %candidate, %134 : i64
%139 = arith.cmpi slt, %candidate, %135 : i64
%140 = arith.andi %138, %139 : i1
cf.assert %140, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%141 = tensor.empty() : tensor<1x4x4xf32>
%142 = 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(%137 : tensor<1x4xi64>) outs(%141 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%143 = arith.index_cast %idx : i64 to index
%144 = linalg.index 2 : index
%145 = tensor.extract %126[%143, %144] : tensor<8x4xf32>
linalg.yield %145 : f32
} -> tensor<1x4x4xf32>
%146 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%147 = tensor.reshape %129(%146) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%148 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%149 = tensor.reshape %130(%148) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%150 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%151 = tensor.reshape %131(%150) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%152 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%153 = tensor.reshape %132(%152) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%154 = arith.constant dense<[1, 4, 8]> : tensor<3xi32>
%155 = tensor.reshape %133(%154) : (tensor<4x8xf32>, tensor<3xi32>) -> tensor<1x4x8xf32>
%156 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%157 = linalg.batch_matmul ins(%142, %147 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%156 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%158 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%159 = linalg.batch_matmul ins(%142, %149 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%158 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%160 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%161 = linalg.batch_matmul ins(%142, %151 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%160 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%162 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%163 = tensor.reshape %157(%162) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%164 = tensor.empty() : tensor<1x2x4x2xf32>
%165 = linalg.transpose ins(%163 : tensor<1x4x2x2xf32>) outs(%164 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%166 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%167 = tensor.reshape %159(%166) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%168 = tensor.empty() : tensor<1x2x2x4xf32>
%169 = linalg.transpose ins(%167 : tensor<1x4x2x2xf32>) outs(%168 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%170 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%171 = tensor.reshape %161(%170) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%172 = tensor.empty() : tensor<1x2x4x2xf32>
%173 = linalg.transpose ins(%171 : tensor<1x4x2x2xf32>) outs(%172 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%174 = tensor.collapse_shape %165 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%175 = tensor.collapse_shape %169 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%176 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%177 = linalg.batch_matmul ins(%174, %175 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%176 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%178 = tensor.expand_shape %177 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%179 = tensor.empty() : tensor<1x2x4x4xf32>
%180 = 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(%128 : tensor<1x1x1x1xf32>) outs(%179 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%181 = tensor.empty() : tensor<1x2x4x4xf32>
%182 = 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(%127 : tensor<1x1x4x4xf32>) outs(%181 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%183 = tensor.empty() : tensor<1x2x4x4xf32>
%184 = linalg.mul ins(%178, %180 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%183 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%185 = tensor.empty() : tensor<1x2x4x4xf32>
%186 = linalg.add ins(%184, %182 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%185 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%189 = arith.constant 0xFF800000 : f32
%187 = tensor.empty() : tensor<1x2x4x1xf32>
%188 = linalg.fill ins(%189 : f32) outs(%187 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%190 = 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(%186 : tensor<1x2x4x4xf32>) outs(%188 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%191 = arith.maximumf %a0, %a1 : f32
linalg.yield %191 : f32
} -> tensor<1x2x4x1xf32>
%192 = tensor.empty() : tensor<1x2x4x4xf32>
%193 = 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(%190 : tensor<1x2x4x1xf32>) outs(%192 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%194 = tensor.empty() : tensor<1x2x4x4xf32>
%195 = linalg.sub ins(%186, %193 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%194 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%196 = tensor.empty() : tensor<1x2x4x4xf32>
%197 = linalg.exp ins(%195 : tensor<1x2x4x4xf32>) outs(%196 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%200 = arith.constant 0xFF800000 : f32
%198 = tensor.empty() : tensor<1x2x4x1xf32>
%199 = linalg.fill ins(%200 : f32) outs(%198 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%201 = 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(%186 : tensor<1x2x4x4xf32>) outs(%199 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%202 = arith.maximumf %a0, %a1 : f32
linalg.yield %202 : f32
} -> tensor<1x2x4x1xf32>
%203 = tensor.empty() : tensor<1x2x4x4xf32>
%204 = 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(%201 : tensor<1x2x4x1xf32>) outs(%203 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%205 = tensor.empty() : tensor<1x2x4x4xf32>
%206 = linalg.sub ins(%186, %204 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%205 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%207 = tensor.empty() : tensor<1x2x4x4xf32>
%208 = linalg.exp ins(%206 : tensor<1x2x4x4xf32>) outs(%207 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%211 = arith.constant 0.000000e+00 : f32
%209 = tensor.empty() : tensor<1x2x4x1xf32>
%210 = linalg.fill ins(%211 : f32) outs(%209 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%212 = 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(%208 : tensor<1x2x4x4xf32>) outs(%210 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%213 = arith.addf %a0, %a1 : f32
linalg.yield %213 : f32
} -> tensor<1x2x4x1xf32>
%214 = tensor.empty() : tensor<1x2x4x4xf32>
%215 = 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(%212 : tensor<1x2x4x1xf32>) outs(%214 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%216 = tensor.empty() : tensor<1x2x4x4xf32>
%217 = linalg.div ins(%197, %215 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%216 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%218 = tensor.collapse_shape %217 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%219 = tensor.collapse_shape %173 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%220 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%221 = linalg.batch_matmul ins(%218, %219 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%220 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%222 = tensor.expand_shape %221 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%223 = tensor.empty() : tensor<1x4x2x2xf32>
%224 = linalg.transpose ins(%222 : tensor<1x2x4x2xf32>) outs(%223 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%225 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%226 = tensor.reshape %224(%225) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%227 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%228 = linalg.batch_matmul ins(%226, %153 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%227 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%229 = tensor.empty() : tensor<1x4x4xf32>
%230 = linalg.add ins(%142, %228 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%229 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%231 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%232 = linalg.batch_matmul ins(%230, %155 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%231 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
return %232 : tensor<1x4x8xf32>
}
func.func @fn_val_233(%234: tensor<1x4xi64>, %235: tensor<8x4xf32>, %236: tensor<1x1x4x4xf32>, %237: tensor<1x1x1x1xf32>, %238: tensor<4x4xf32>, %239: tensor<4x4xf32>, %240: tensor<4x4xf32>, %241: tensor<4x4xf32>, %242: tensor<4x8xf32>, %243: tensor<1x1xf32>) -> tensor<8x4xf32> attributes {llvm.emit_c_interface} {
%244 = arith.constant 0 : i64
%245 = arith.constant 8 : i64
%246 = tensor.empty() : tensor<1x4xi64>
%247 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%234 : tensor<1x4xi64>) outs(%246 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%248 = arith.cmpi sge, %candidate, %244 : i64
%249 = arith.cmpi slt, %candidate, %245 : i64
%250 = arith.andi %248, %249 : i1
cf.assert %250, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%251 = tensor.empty() : tensor<1x4x4xf32>
%252 = 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(%247 : tensor<1x4xi64>) outs(%251 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%253 = arith.index_cast %idx : i64 to index
%254 = linalg.index 2 : index
%255 = tensor.extract %235[%253, %254] : tensor<8x4xf32>
linalg.yield %255 : f32
} -> tensor<1x4x4xf32>
%256 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%257 = tensor.reshape %238(%256) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%258 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%259 = tensor.reshape %239(%258) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%260 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%261 = tensor.reshape %240(%260) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%262 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%263 = tensor.reshape %241(%262) : (tensor<4x4xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%264 = arith.constant dense<[1, 4, 8]> : tensor<3xi32>
%265 = tensor.reshape %242(%264) : (tensor<4x8xf32>, tensor<3xi32>) -> tensor<1x4x8xf32>
%266 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%267 = linalg.batch_matmul ins(%252, %257 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%266 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%268 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%269 = linalg.batch_matmul ins(%252, %259 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%268 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%270 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%271 = linalg.batch_matmul ins(%252, %261 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%270 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%272 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%273 = tensor.reshape %267(%272) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%274 = tensor.empty() : tensor<1x2x4x2xf32>
%275 = linalg.transpose ins(%273 : tensor<1x4x2x2xf32>) outs(%274 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%276 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%277 = tensor.reshape %269(%276) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%278 = tensor.empty() : tensor<1x2x2x4xf32>
%279 = linalg.transpose ins(%277 : tensor<1x4x2x2xf32>) outs(%278 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%280 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%281 = tensor.reshape %271(%280) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%282 = tensor.empty() : tensor<1x2x4x2xf32>
%283 = linalg.transpose ins(%281 : tensor<1x4x2x2xf32>) outs(%282 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%284 = tensor.collapse_shape %275 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%285 = tensor.collapse_shape %279 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%286 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%287 = linalg.batch_matmul ins(%284, %285 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%286 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%288 = tensor.expand_shape %287 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%289 = tensor.empty() : tensor<1x2x4x4xf32>
%290 = 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(%237 : tensor<1x1x1x1xf32>) outs(%289 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%291 = tensor.empty() : tensor<1x2x4x4xf32>
%292 = 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(%236 : tensor<1x1x4x4xf32>) outs(%291 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%293 = tensor.empty() : tensor<1x2x4x4xf32>
%294 = linalg.mul ins(%288, %290 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%293 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%295 = tensor.empty() : tensor<1x2x4x4xf32>
%296 = linalg.add ins(%294, %292 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%295 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%299 = arith.constant 0xFF800000 : f32
%297 = tensor.empty() : tensor<1x2x4x1xf32>
%298 = linalg.fill ins(%299 : f32) outs(%297 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%300 = 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(%296 : tensor<1x2x4x4xf32>) outs(%298 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%301 = arith.maximumf %a0, %a1 : f32
linalg.yield %301 : f32
} -> tensor<1x2x4x1xf32>
%302 = tensor.empty() : tensor<1x2x4x4xf32>
%303 = 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(%300 : tensor<1x2x4x1xf32>) outs(%302 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%304 = tensor.empty() : tensor<1x2x4x4xf32>
%305 = linalg.sub ins(%296, %303 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%304 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%306 = tensor.empty() : tensor<1x2x4x4xf32>
%307 = linalg.exp ins(%305 : tensor<1x2x4x4xf32>) outs(%306 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%310 = arith.constant 0xFF800000 : f32
%308 = tensor.empty() : tensor<1x2x4x1xf32>
%309 = linalg.fill ins(%310 : f32) outs(%308 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%311 = 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(%296 : tensor<1x2x4x4xf32>) outs(%309 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%312 = arith.maximumf %a0, %a1 : f32
linalg.yield %312 : f32
} -> tensor<1x2x4x1xf32>
%313 = tensor.empty() : tensor<1x2x4x4xf32>
%314 = 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(%311 : tensor<1x2x4x1xf32>) outs(%313 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%315 = tensor.empty() : tensor<1x2x4x4xf32>
%316 = linalg.sub ins(%296, %314 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%315 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%317 = tensor.empty() : tensor<1x2x4x4xf32>
%318 = linalg.exp ins(%316 : tensor<1x2x4x4xf32>) outs(%317 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%321 = arith.constant 0.000000e+00 : f32
%319 = tensor.empty() : tensor<1x2x4x1xf32>
%320 = linalg.fill ins(%321 : f32) outs(%319 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%322 = 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(%318 : tensor<1x2x4x4xf32>) outs(%320 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%323 = arith.addf %a0, %a1 : f32
linalg.yield %323 : f32
} -> tensor<1x2x4x1xf32>
%324 = tensor.empty() : tensor<1x2x4x4xf32>
%325 = 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(%322 : tensor<1x2x4x1xf32>) outs(%324 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%326 = tensor.empty() : tensor<1x2x4x4xf32>
%327 = linalg.div ins(%307, %325 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%326 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%328 = tensor.collapse_shape %327 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%329 = tensor.collapse_shape %283 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%330 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%331 = linalg.batch_matmul ins(%328, %329 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%330 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%332 = tensor.expand_shape %331 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%333 = tensor.empty() : tensor<1x4x2x2xf32>
%334 = linalg.transpose ins(%332 : tensor<1x2x4x2xf32>) outs(%333 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%335 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%336 = tensor.reshape %334(%335) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%337 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%338 = linalg.batch_matmul ins(%336, %263 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%337 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%339 = tensor.empty() : tensor<1x4x4xf32>
%340 = linalg.add ins(%252, %338 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%339 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%341 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%342 = linalg.batch_matmul ins(%340, %265 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%341 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%343 = arith.constant dense<1.000000e+00> : tensor<1x4x8xf32>
%344 = tensor.empty() : tensor<1x4x4xf32>
%345 = linalg.transpose ins(%340 : tensor<1x4x4xf32>) outs(%344 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%346 = arith.constant dense<1.000000e+00> : tensor<1x4x8xf32>
%347 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%348 = linalg.batch_matmul ins(%345, %346 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%347 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
%349 = tensor.empty() : tensor<1x8x4xf32>
%350 = linalg.transpose ins(%265 : tensor<1x4x8xf32>) outs(%349 : tensor<1x8x4xf32>) permutation = [0, 2, 1]
%351 = arith.constant dense<1.000000e+00> : tensor<1x4x8xf32>
%352 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%353 = linalg.batch_matmul ins(%351, %350 : tensor<1x4x8xf32>, tensor<1x8x4xf32>) outs(%352 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%354 = tensor.empty() : tensor<1x4x4xf32>
%355 = linalg.transpose ins(%336 : tensor<1x4x4xf32>) outs(%354 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%356 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%357 = linalg.batch_matmul ins(%355, %353 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%356 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%358 = tensor.empty() : tensor<1x4x4xf32>
%359 = linalg.transpose ins(%263 : tensor<1x4x4xf32>) outs(%358 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%360 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%361 = linalg.batch_matmul ins(%353, %359 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%360 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%362 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%363 = tensor.reshape %361(%362) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%364 = tensor.empty() : tensor<1x2x4x2xf32>
%365 = linalg.transpose ins(%363 : tensor<1x4x2x2xf32>) outs(%364 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%366 = tensor.empty() : tensor<1x2x4x4xf32>
%367 = linalg.transpose ins(%327 : tensor<1x2x4x4xf32>) outs(%366 : tensor<1x2x4x4xf32>) permutation = [0, 1, 3, 2]
%368 = tensor.collapse_shape %367 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%369 = tensor.collapse_shape %365 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%370 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%371 = linalg.batch_matmul ins(%368, %369 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%370 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%372 = tensor.expand_shape %371 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%373 = tensor.empty() : tensor<1x2x2x4xf32>
%374 = linalg.transpose ins(%283 : tensor<1x2x4x2xf32>) outs(%373 : tensor<1x2x2x4xf32>) permutation = [0, 1, 3, 2]
%375 = tensor.collapse_shape %365 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%376 = tensor.collapse_shape %374 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%377 = arith.constant dense<0.000000e+00> : tensor<2x4x4xf32>
%378 = linalg.batch_matmul ins(%375, %376 : tensor<2x4x2xf32>, tensor<2x2x4xf32>) outs(%377 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32>
%379 = tensor.expand_shape %378 [[0, 1], [2], [3]] : tensor<2x4x4xf32> into tensor<1x2x4x4xf32>
%380 = tensor.empty() : tensor<1x2x4x4xf32>
%381 = linalg.mul ins(%327, %379 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%380 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%384 = arith.constant 0.000000e+00 : f32
%382 = tensor.empty() : tensor<1x2x4x1xf32>
%383 = linalg.fill ins(%384 : f32) outs(%382 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%385 = 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(%381 : tensor<1x2x4x4xf32>) outs(%383 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%386 = arith.addf %a0, %a1 : f32
linalg.yield %386 : f32
} -> tensor<1x2x4x1xf32>
%387 = tensor.empty() : tensor<1x2x4x4xf32>
%388 = 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(%385 : tensor<1x2x4x1xf32>) outs(%387 : tensor<1x2x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x2x4x4xf32>
%389 = tensor.empty() : tensor<1x2x4x4xf32>
%390 = linalg.sub ins(%379, %388 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%389 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%391 = tensor.empty() : tensor<1x2x4x4xf32>
%392 = linalg.mul ins(%327, %390 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%391 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%393 = tensor.empty() : tensor<1x2x4x4xf32>
%394 = linalg.mul ins(%392, %290 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%393 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%395 = tensor.empty() : tensor<1x2x4x4xf32>
%396 = linalg.mul ins(%392, %288 : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) outs(%395 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32>
%399 = arith.constant 0.000000e+00 : f32
%397 = tensor.empty() : tensor<1x1x4x4xf32>
%398 = linalg.fill ins(%399 : f32) outs(%397 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%400 = 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(%392 : tensor<1x2x4x4xf32>) outs(%398 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
%401 = arith.addf %a0, %a1 : f32
linalg.yield %401 : f32
} -> tensor<1x1x4x4xf32>
%404 = arith.constant 0.000000e+00 : f32
%402 = tensor.empty() : tensor<1x2x4x1xf32>
%403 = linalg.fill ins(%404 : f32) outs(%402 : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
%405 = 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(%396 : tensor<1x2x4x4xf32>) outs(%403 : tensor<1x2x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%406 = arith.addf %a0, %a1 : f32
linalg.yield %406 : f32
} -> tensor<1x2x4x1xf32>
%409 = arith.constant 0.000000e+00 : f32
%407 = tensor.empty() : tensor<1x2x1x1xf32>
%408 = linalg.fill ins(%409 : f32) outs(%407 : tensor<1x2x1x1xf32>) -> tensor<1x2x1x1xf32>
%410 = 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(%405 : tensor<1x2x4x1xf32>) outs(%408 : tensor<1x2x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%411 = arith.addf %a0, %a1 : f32
linalg.yield %411 : f32
} -> tensor<1x2x1x1xf32>
%414 = arith.constant 0.000000e+00 : f32
%412 = tensor.empty() : tensor<1x1x1x1xf32>
%413 = linalg.fill ins(%414 : f32) outs(%412 : tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
%415 = 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(%410 : tensor<1x2x1x1xf32>) outs(%413 : tensor<1x1x1x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%416 = arith.addf %a0, %a1 : f32
linalg.yield %416 : f32
} -> tensor<1x1x1x1xf32>
%417 = tensor.empty() : tensor<1x2x2x4xf32>
%418 = linalg.transpose ins(%275 : tensor<1x2x4x2xf32>) outs(%417 : tensor<1x2x2x4xf32>) permutation = [0, 1, 3, 2]
%419 = tensor.collapse_shape %418 [[0, 1], [2], [3]] : tensor<1x2x2x4xf32> into tensor<2x2x4xf32>
%420 = tensor.collapse_shape %394 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%421 = arith.constant dense<0.000000e+00> : tensor<2x2x4xf32>
%422 = linalg.batch_matmul ins(%419, %420 : tensor<2x2x4xf32>, tensor<2x4x4xf32>) outs(%421 : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
%423 = tensor.expand_shape %422 [[0, 1], [2], [3]] : tensor<2x2x4xf32> into tensor<1x2x2x4xf32>
%424 = tensor.empty() : tensor<1x2x4x2xf32>
%425 = linalg.transpose ins(%279 : tensor<1x2x2x4xf32>) outs(%424 : tensor<1x2x4x2xf32>) permutation = [0, 1, 3, 2]
%426 = tensor.collapse_shape %394 [[0, 1], [2], [3]] : tensor<1x2x4x4xf32> into tensor<2x4x4xf32>
%427 = tensor.collapse_shape %425 [[0, 1], [2], [3]] : tensor<1x2x4x2xf32> into tensor<2x4x2xf32>
%428 = arith.constant dense<0.000000e+00> : tensor<2x4x2xf32>
%429 = linalg.batch_matmul ins(%426, %427 : tensor<2x4x4xf32>, tensor<2x4x2xf32>) outs(%428 : tensor<2x4x2xf32>) -> tensor<2x4x2xf32>
%430 = tensor.expand_shape %429 [[0, 1], [2], [3]] : tensor<2x4x2xf32> into tensor<1x2x4x2xf32>
%431 = tensor.empty() : tensor<1x4x2x2xf32>
%432 = linalg.transpose ins(%372 : tensor<1x2x4x2xf32>) outs(%431 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%433 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%434 = tensor.reshape %432(%433) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%435 = tensor.empty() : tensor<1x4x2x2xf32>
%436 = linalg.transpose ins(%423 : tensor<1x2x2x4xf32>) outs(%435 : tensor<1x4x2x2xf32>) permutation = [0, 3, 1, 2]
%437 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%438 = tensor.reshape %436(%437) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%439 = tensor.empty() : tensor<1x4x2x2xf32>
%440 = linalg.transpose ins(%430 : tensor<1x2x4x2xf32>) outs(%439 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%441 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%442 = tensor.reshape %440(%441) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%443 = tensor.empty() : tensor<1x4x4xf32>
%444 = linalg.transpose ins(%252 : tensor<1x4x4xf32>) outs(%443 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%445 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%446 = linalg.batch_matmul ins(%444, %434 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%445 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%447 = tensor.empty() : tensor<1x4x4xf32>
%448 = linalg.transpose ins(%261 : tensor<1x4x4xf32>) outs(%447 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%449 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%450 = linalg.batch_matmul ins(%434, %448 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%449 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%451 = tensor.empty() : tensor<1x4x4xf32>
%452 = linalg.transpose ins(%252 : tensor<1x4x4xf32>) outs(%451 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%453 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%454 = linalg.batch_matmul ins(%452, %438 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%453 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%455 = tensor.empty() : tensor<1x4x4xf32>
%456 = linalg.transpose ins(%259 : tensor<1x4x4xf32>) outs(%455 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%457 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%458 = linalg.batch_matmul ins(%438, %456 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%457 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%459 = tensor.empty() : tensor<1x4x4xf32>
%460 = linalg.transpose ins(%252 : tensor<1x4x4xf32>) outs(%459 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%461 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%462 = linalg.batch_matmul ins(%460, %442 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%461 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%463 = tensor.empty() : tensor<1x4x4xf32>
%464 = linalg.transpose ins(%257 : tensor<1x4x4xf32>) outs(%463 : tensor<1x4x4xf32>) permutation = [0, 2, 1]
%465 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%466 = linalg.batch_matmul ins(%442, %464 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%465 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%467 = arith.constant dense<[4, 8]> : tensor<2xi32>
%468 = tensor.reshape %348(%467) : (tensor<1x4x8xf32>, tensor<2xi32>) -> tensor<4x8xf32>
%469 = arith.constant dense<[4, 4]> : tensor<2xi32>
%470 = tensor.reshape %357(%469) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%471 = arith.constant dense<[4, 4]> : tensor<2xi32>
%472 = tensor.reshape %446(%471) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%473 = arith.constant dense<[4, 4]> : tensor<2xi32>
%474 = tensor.reshape %454(%473) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%475 = arith.constant dense<[4, 4]> : tensor<2xi32>
%476 = tensor.reshape %462(%475) : (tensor<1x4x4xf32>, tensor<2xi32>) -> tensor<4x4xf32>
%477 = tensor.empty() : tensor<1x4x4xf32>
%478 = linalg.add ins(%353, %450 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%477 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%479 = tensor.empty() : tensor<1x4x4xf32>
%480 = linalg.add ins(%478, %458 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%479 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%481 = tensor.empty() : tensor<1x4x4xf32>
%482 = linalg.add ins(%480, %466 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%481 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%483 = arith.constant 0 : i64
%484 = arith.constant 8 : i64
%485 = tensor.empty() : tensor<1x4xi64>
%486 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%234 : tensor<1x4xi64>) outs(%485 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%487 = arith.cmpi sge, %candidate, %483 : i64
%488 = arith.cmpi slt, %candidate, %484 : i64
%489 = arith.andi %487, %488 : i1
cf.assert %489, "scatter_add index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%490 = tensor.collapse_shape %486 [[0, 1]] : tensor<1x4xi64> into tensor<4xi64>
%491 = tensor.collapse_shape %482 [[0, 1], [2]] : tensor<1x4x4xf32> into tensor<4x4xf32>
%492 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%493 = 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(%490, %491 : tensor<4xi64>, tensor<4x4xf32>) outs(%492 : tensor<8x4xf32>) {
^bb0(%idx: i64, %upd: f32, %acc: f32):
%494 = linalg.index 0 : index
%495 = arith.index_cast %494 : index to i64
%496 = arith.cmpi eq, %idx, %495 : i64
%497 = arith.constant 0.000000e+00 : f32
%498 = arith.select %496, %upd, %497 : f32
%499 = arith.addf %acc, %498 : f32
linalg.yield %499 : f32
} -> tensor<8x4xf32>
%500 = tensor.empty() : tensor<8x4xf32>
%501 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%243 : tensor<1x1xf32>) outs(%500 : tensor<8x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<8x4xf32>
%502 = tensor.empty() : tensor<8x4xf32>
%503 = linalg.mul ins(%501, %493 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%502 : tensor<8x4xf32>) -> tensor<8x4xf32>
%504 = tensor.empty() : tensor<8x4xf32>
%505 = linalg.sub ins(%235, %503 : tensor<8x4xf32>, tensor<8x4xf32>) outs(%504 : tensor<8x4xf32>) -> tensor<8x4xf32>
return %505 : tensor<8x4xf32>
}
func.func @main() -> tensor<8x4xf32> {
%506 = arith.constant dense<0> : tensor<1x4xi64>
%507 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%508 = arith.constant dense<0.000000e+00> : tensor<1x1x4x4xf32>
%509 = arith.constant dense<0.000000e+00> : tensor<1x1x1x1xf32>
%510 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%511 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%512 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%513 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%514 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%515 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%516 = func.call @fn_val_233(%506, %507, %508, %509, %510, %511, %512, %513, %514, %515) : (tensor<1x4xi64>, tensor<8x4xf32>, tensor<1x1x4x4xf32>, tensor<1x1x1x1xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>, tensor<1x1xf32>) -> tensor<8x4xf32>
return %516 : tensor<8x4xf32>
}
}