Example
storylm0_tp_heads.kina
let tp_heads = fn(
q: tensor<F32, [1, 2, 4, 2], @host>,
k_t: tensor<F32, [1, 2, 2, 4], @host>,
v: tensor<F32, [1, 2, 4, 2], @host>,
attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
causal_mask: tensor<F32, [1, 1, 4, 4], @host>
) -> tensor<F32, [1, 2, 4, 2], @host> {
let q0 = q.slice(axis=1, start=0, stop=1) in
let k0 = k_t.slice(axis=1, start=0, stop=1) in
let v0 = v.slice(axis=1, start=0, stop=1) in
let q1 = q.slice(axis=1, start=1, stop=2) in
let k1 = k_t.slice(axis=1, start=1, stop=2) in
let v1 = v.slice(axis=1, start=1, stop=2) in
let scale_b = attention_scale.broadcast([1, 1, 4, 4]) in
let mask_b = causal_mask.broadcast([1, 1, 4, 4]) in
let scores0 = q0 @ k0 in
let scores1 = q1 @ k1 in
let probs0 = ((scores0 * scale_b) + mask_b).softmax(axis=3) in
let probs1 = ((scores1 * scale_b) + mask_b).softmax(axis=3) in
let ctx0 = probs0 @ v0 in
let ctx1 = probs1 @ v1 in
let context = ctx0.concat(axis=1, ctx1) in
context
} in
let storylm0_tp_heads = 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
let wq_broadcast = wq.broadcast([1, 4, 4]) in
let wk_broadcast = wk.broadcast([1, 4, 4]) in
let wv_broadcast = wv.broadcast([1, 4, 4]) in
let wo_broadcast = wo.broadcast([1, 4, 4]) in
let lm_head_broadcast = lm_head.broadcast([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 = q_flat.reshape([1, 4, 2, 2]).transpose([0, 2, 1, 3]) in
let k_t = k_flat.reshape([1, 4, 2, 2]).transpose([0, 2, 3, 1]) in
let v = v_flat.reshape([1, 4, 2, 2]).transpose([0, 2, 1, 3]) in
let context = tp_heads(q, k_t, v, attention_scale, causal_mask) in
let merged = context.transpose([0, 2, 1, 3]).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 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
storylm0_tp_heads(token_ids, token_embedding, causal_mask, attention_scale, wq, wk, wv, wo, lm_head)
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<1x2x4x2xf32>, %126: tensor<1x2x2x4xf32>, %127: tensor<1x2x4x2xf32>, %128: tensor<1x1x1x1xf32>, %129: tensor<1x1x4x4xf32>) -> tensor<1x2x4x2xf32> attributes {llvm.emit_c_interface} {
%130 = tensor.extract_slice %125[0, 0, 0, 0][1, 1, 4, 2][1, 1, 1, 1] : tensor<1x2x4x2xf32> to tensor<1x1x4x2xf32>
%131 = tensor.extract_slice %126[0, 0, 0, 0][1, 1, 2, 4][1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x1x2x4xf32>
%132 = tensor.extract_slice %127[0, 0, 0, 0][1, 1, 4, 2][1, 1, 1, 1] : tensor<1x2x4x2xf32> to tensor<1x1x4x2xf32>
%133 = tensor.extract_slice %125[0, 1, 0, 0][1, 1, 4, 2][1, 1, 1, 1] : tensor<1x2x4x2xf32> to tensor<1x1x4x2xf32>
%134 = tensor.extract_slice %126[0, 1, 0, 0][1, 1, 2, 4][1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x1x2x4xf32>
%135 = tensor.extract_slice %127[0, 1, 0, 0][1, 1, 4, 2][1, 1, 1, 1] : tensor<1x2x4x2xf32> to tensor<1x1x4x2xf32>
%136 = tensor.empty() : tensor<1x1x4x4xf32>
%137 = 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(%136 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%138 = tensor.empty() : tensor<1x1x4x4xf32>
%139 = 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(%129 : tensor<1x1x4x4xf32>) outs(%138 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%140 = tensor.collapse_shape %130 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
%141 = tensor.collapse_shape %131 [[0, 1], [2], [3]] : tensor<1x1x2x4xf32> into tensor<1x2x4xf32>
%142 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%143 = linalg.batch_matmul ins(%140, %141 : tensor<1x4x2xf32>, tensor<1x2x4xf32>) outs(%142 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%144 = tensor.expand_shape %143 [[0, 1], [2], [3]] : tensor<1x4x4xf32> into tensor<1x1x4x4xf32>
%145 = tensor.collapse_shape %133 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
%146 = tensor.collapse_shape %134 [[0, 1], [2], [3]] : tensor<1x1x2x4xf32> into tensor<1x2x4xf32>
%147 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%148 = linalg.batch_matmul ins(%145, %146 : tensor<1x4x2xf32>, tensor<1x2x4xf32>) outs(%147 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%149 = tensor.expand_shape %148 [[0, 1], [2], [3]] : tensor<1x4x4xf32> into tensor<1x1x4x4xf32>
%150 = tensor.empty() : tensor<1x1x4x4xf32>
%151 = linalg.mul ins(%144, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%150 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%152 = tensor.empty() : tensor<1x1x4x4xf32>
%153 = linalg.add ins(%151, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%152 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%154 = tensor.empty() : tensor<1x1x4x4xf32>
%155 = linalg.mul ins(%144, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%154 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%156 = tensor.empty() : tensor<1x1x4x4xf32>
%157 = linalg.add ins(%155, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%156 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%160 = arith.constant 0xFF800000 : f32
%158 = tensor.empty() : tensor<1x1x4x1xf32>
%159 = linalg.fill ins(%160 : f32) outs(%158 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
%161 = 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(%157 : tensor<1x1x4x4xf32>) outs(%159 : tensor<1x1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%162 = arith.maximumf %a0, %a1 : f32
linalg.yield %162 : f32
} -> tensor<1x1x4x1xf32>
%163 = tensor.empty() : tensor<1x1x4x4xf32>
%164 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%161 : tensor<1x1x4x1xf32>) outs(%163 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%165 = tensor.empty() : tensor<1x1x4x4xf32>
%166 = linalg.sub ins(%153, %164 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%165 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%167 = tensor.empty() : tensor<1x1x4x4xf32>
%168 = linalg.exp ins(%166 : tensor<1x1x4x4xf32>) outs(%167 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%169 = tensor.empty() : tensor<1x1x4x4xf32>
%170 = linalg.mul ins(%144, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%169 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%171 = tensor.empty() : tensor<1x1x4x4xf32>
%172 = linalg.add ins(%170, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%171 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%173 = tensor.empty() : tensor<1x1x4x4xf32>
%174 = linalg.mul ins(%144, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%173 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%175 = tensor.empty() : tensor<1x1x4x4xf32>
%176 = linalg.add ins(%174, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%175 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%179 = arith.constant 0xFF800000 : f32
%177 = tensor.empty() : tensor<1x1x4x1xf32>
%178 = linalg.fill ins(%179 : f32) outs(%177 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
%180 = 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(%176 : tensor<1x1x4x4xf32>) outs(%178 : tensor<1x1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%181 = arith.maximumf %a0, %a1 : f32
linalg.yield %181 : f32
} -> tensor<1x1x4x1xf32>
%182 = tensor.empty() : tensor<1x1x4x4xf32>
%183 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%180 : tensor<1x1x4x1xf32>) outs(%182 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%184 = tensor.empty() : tensor<1x1x4x4xf32>
%185 = linalg.sub ins(%172, %183 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%184 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%186 = tensor.empty() : tensor<1x1x4x4xf32>
%187 = linalg.exp ins(%185 : tensor<1x1x4x4xf32>) outs(%186 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%190 = arith.constant 0.000000e+00 : f32
%188 = tensor.empty() : tensor<1x1x4x1xf32>
%189 = linalg.fill ins(%190 : f32) outs(%188 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
%191 = 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(%187 : tensor<1x1x4x4xf32>) outs(%189 : tensor<1x1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%192 = arith.addf %a0, %a1 : f32
linalg.yield %192 : f32
} -> tensor<1x1x4x1xf32>
%193 = tensor.empty() : tensor<1x1x4x4xf32>
%194 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%191 : tensor<1x1x4x1xf32>) outs(%193 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%195 = tensor.empty() : tensor<1x1x4x4xf32>
%196 = linalg.div ins(%168, %194 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%195 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%197 = tensor.empty() : tensor<1x1x4x4xf32>
%198 = linalg.mul ins(%149, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%197 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%199 = tensor.empty() : tensor<1x1x4x4xf32>
%200 = linalg.add ins(%198, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%199 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%201 = tensor.empty() : tensor<1x1x4x4xf32>
%202 = linalg.mul ins(%149, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%201 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%203 = tensor.empty() : tensor<1x1x4x4xf32>
%204 = linalg.add ins(%202, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%203 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%207 = arith.constant 0xFF800000 : f32
%205 = tensor.empty() : tensor<1x1x4x1xf32>
%206 = linalg.fill ins(%207 : f32) outs(%205 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
%208 = 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(%204 : tensor<1x1x4x4xf32>) outs(%206 : tensor<1x1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%209 = arith.maximumf %a0, %a1 : f32
linalg.yield %209 : f32
} -> tensor<1x1x4x1xf32>
%210 = tensor.empty() : tensor<1x1x4x4xf32>
%211 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%208 : tensor<1x1x4x1xf32>) outs(%210 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%212 = tensor.empty() : tensor<1x1x4x4xf32>
%213 = linalg.sub ins(%200, %211 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%212 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%214 = tensor.empty() : tensor<1x1x4x4xf32>
%215 = linalg.exp ins(%213 : tensor<1x1x4x4xf32>) outs(%214 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%216 = tensor.empty() : tensor<1x1x4x4xf32>
%217 = linalg.mul ins(%149, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%216 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%218 = tensor.empty() : tensor<1x1x4x4xf32>
%219 = linalg.add ins(%217, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%218 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%220 = tensor.empty() : tensor<1x1x4x4xf32>
%221 = linalg.mul ins(%149, %137 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%220 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%222 = tensor.empty() : tensor<1x1x4x4xf32>
%223 = linalg.add ins(%221, %139 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%222 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%226 = arith.constant 0xFF800000 : f32
%224 = tensor.empty() : tensor<1x1x4x1xf32>
%225 = linalg.fill ins(%226 : f32) outs(%224 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
%227 = 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(%223 : tensor<1x1x4x4xf32>) outs(%225 : tensor<1x1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%228 = arith.maximumf %a0, %a1 : f32
linalg.yield %228 : f32
} -> tensor<1x1x4x1xf32>
%229 = tensor.empty() : tensor<1x1x4x4xf32>
%230 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%227 : tensor<1x1x4x1xf32>) outs(%229 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%231 = tensor.empty() : tensor<1x1x4x4xf32>
%232 = linalg.sub ins(%219, %230 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%231 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%233 = tensor.empty() : tensor<1x1x4x4xf32>
%234 = linalg.exp ins(%232 : tensor<1x1x4x4xf32>) outs(%233 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%237 = arith.constant 0.000000e+00 : f32
%235 = tensor.empty() : tensor<1x1x4x1xf32>
%236 = linalg.fill ins(%237 : f32) outs(%235 : tensor<1x1x4x1xf32>) -> tensor<1x1x4x1xf32>
%238 = 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(%234 : tensor<1x1x4x4xf32>) outs(%236 : tensor<1x1x4x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%239 = arith.addf %a0, %a1 : f32
linalg.yield %239 : f32
} -> tensor<1x1x4x1xf32>
%240 = tensor.empty() : tensor<1x1x4x4xf32>
%241 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (0, 0, d2, 0)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%238 : tensor<1x1x4x1xf32>) outs(%240 : tensor<1x1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x1x4x4xf32>
%242 = tensor.empty() : tensor<1x1x4x4xf32>
%243 = linalg.div ins(%215, %241 : tensor<1x1x4x4xf32>, tensor<1x1x4x4xf32>) outs(%242 : tensor<1x1x4x4xf32>) -> tensor<1x1x4x4xf32>
%244 = tensor.collapse_shape %196 [[0, 1], [2], [3]] : tensor<1x1x4x4xf32> into tensor<1x4x4xf32>
%245 = tensor.collapse_shape %132 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
%246 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
%247 = linalg.batch_matmul ins(%244, %245 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%246 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
%248 = tensor.expand_shape %247 [[0, 1], [2], [3]] : tensor<1x4x2xf32> into tensor<1x1x4x2xf32>
%249 = tensor.collapse_shape %243 [[0, 1], [2], [3]] : tensor<1x1x4x4xf32> into tensor<1x4x4xf32>
%250 = tensor.collapse_shape %135 [[0, 1], [2], [3]] : tensor<1x1x4x2xf32> into tensor<1x4x2xf32>
%251 = arith.constant dense<0.000000e+00> : tensor<1x4x2xf32>
%252 = linalg.batch_matmul ins(%249, %250 : tensor<1x4x4xf32>, tensor<1x4x2xf32>) outs(%251 : tensor<1x4x2xf32>) -> tensor<1x4x2xf32>
%253 = tensor.expand_shape %252 [[0, 1], [2], [3]] : tensor<1x4x2xf32> into tensor<1x1x4x2xf32>
%254 = tensor.empty() : tensor<1x2x4x2xf32>
%255 = tensor.insert_slice %248 into %254[0, 0, 0, 0] [1, 1, 4, 2] [1, 1, 1, 1] : tensor<1x1x4x2xf32> into tensor<1x2x4x2xf32>
%256 = tensor.insert_slice %253 into %255[0, 1, 0, 0] [1, 1, 4, 2] [1, 1, 1, 1] : tensor<1x1x4x2xf32> into tensor<1x2x4x2xf32>
return %256 : tensor<1x2x4x2xf32>
}
func.func @fn_val_257(%258: tensor<1x4xi64>, %259: tensor<8x4xf32>, %260: tensor<1x1x4x4xf32>, %261: tensor<1x1x1x1xf32>, %262: tensor<4x4xf32>, %263: tensor<4x4xf32>, %264: tensor<4x4xf32>, %265: tensor<4x4xf32>, %266: tensor<4x8xf32>) -> tensor<1x4x8xf32> attributes {llvm.emit_c_interface} {
%267 = arith.constant 0 : i64
%268 = arith.constant 8 : i64
%269 = tensor.empty() : tensor<1x4xi64>
%270 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%258 : tensor<1x4xi64>) outs(%269 : tensor<1x4xi64>) {
^bb0(%candidate: i64, %_: i64):
%271 = arith.cmpi sge, %candidate, %267 : i64
%272 = arith.cmpi slt, %candidate, %268 : i64
%273 = arith.andi %271, %272 : i1
cf.assert %273, "gather index out of bounds"
linalg.yield %candidate : i64
} -> tensor<1x4xi64>
%274 = tensor.empty() : tensor<1x4x4xf32>
%275 = 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(%270 : tensor<1x4xi64>) outs(%274 : tensor<1x4x4xf32>) {
^bb0(%idx: i64, %_: f32):
%276 = arith.index_cast %idx : i64 to index
%277 = linalg.index 2 : index
%278 = tensor.extract %259[%276, %277] : tensor<8x4xf32>
linalg.yield %278 : f32
} -> tensor<1x4x4xf32>
%279 = tensor.empty() : tensor<1x4x4xf32>
%280 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%262 : tensor<4x4xf32>) outs(%279 : tensor<1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x4xf32>
%281 = tensor.empty() : tensor<1x4x4xf32>
%282 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%263 : tensor<4x4xf32>) outs(%281 : tensor<1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x4xf32>
%283 = tensor.empty() : tensor<1x4x4xf32>
%284 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%264 : tensor<4x4xf32>) outs(%283 : tensor<1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x4xf32>
%285 = tensor.empty() : tensor<1x4x4xf32>
%286 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%265 : tensor<4x4xf32>) outs(%285 : tensor<1x4x4xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x4xf32>
%287 = tensor.empty() : tensor<1x4x8xf32>
%288 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%266 : tensor<4x8xf32>) outs(%287 : tensor<1x4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x4x8xf32>
%289 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%290 = linalg.batch_matmul ins(%275, %280 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%289 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%291 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%292 = linalg.batch_matmul ins(%275, %282 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%291 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%293 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%294 = linalg.batch_matmul ins(%275, %284 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%293 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%295 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%296 = tensor.reshape %290(%295) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%297 = tensor.empty() : tensor<1x2x4x2xf32>
%298 = linalg.transpose ins(%296 : tensor<1x4x2x2xf32>) outs(%297 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%299 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%300 = tensor.reshape %292(%299) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%301 = tensor.empty() : tensor<1x2x2x4xf32>
%302 = linalg.transpose ins(%300 : tensor<1x4x2x2xf32>) outs(%301 : tensor<1x2x2x4xf32>) permutation = [0, 2, 3, 1]
%303 = arith.constant dense<[1, 4, 2, 2]> : tensor<4xi32>
%304 = tensor.reshape %294(%303) : (tensor<1x4x4xf32>, tensor<4xi32>) -> tensor<1x4x2x2xf32>
%305 = tensor.empty() : tensor<1x2x4x2xf32>
%306 = linalg.transpose ins(%304 : tensor<1x4x2x2xf32>) outs(%305 : tensor<1x2x4x2xf32>) permutation = [0, 2, 1, 3]
%307 = func.call @fn_val_124(%298, %302, %306, %261, %260) : (tensor<1x2x4x2xf32>, tensor<1x2x2x4xf32>, tensor<1x2x4x2xf32>, tensor<1x1x1x1xf32>, tensor<1x1x4x4xf32>) -> tensor<1x2x4x2xf32>
%308 = tensor.empty() : tensor<1x4x2x2xf32>
%309 = linalg.transpose ins(%307 : tensor<1x2x4x2xf32>) outs(%308 : tensor<1x4x2x2xf32>) permutation = [0, 2, 1, 3]
%310 = arith.constant dense<[1, 4, 4]> : tensor<3xi32>
%311 = tensor.reshape %309(%310) : (tensor<1x4x2x2xf32>, tensor<3xi32>) -> tensor<1x4x4xf32>
%312 = arith.constant dense<0.000000e+00> : tensor<1x4x4xf32>
%313 = linalg.batch_matmul ins(%311, %286 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%312 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%314 = tensor.empty() : tensor<1x4x4xf32>
%315 = linalg.add ins(%275, %313 : tensor<1x4x4xf32>, tensor<1x4x4xf32>) outs(%314 : tensor<1x4x4xf32>) -> tensor<1x4x4xf32>
%316 = arith.constant dense<0.000000e+00> : tensor<1x4x8xf32>
%317 = linalg.batch_matmul ins(%315, %288 : tensor<1x4x4xf32>, tensor<1x4x8xf32>) outs(%316 : tensor<1x4x8xf32>) -> tensor<1x4x8xf32>
return %317 : tensor<1x4x8xf32>
}
func.func @main() -> tensor<1x4x8xf32> {
%318 = arith.constant dense<0> : tensor<1x4xi64>
%319 = arith.constant dense<0.000000e+00> : tensor<8x4xf32>
%320 = arith.constant dense<0.000000e+00> : tensor<1x1x4x4xf32>
%321 = arith.constant dense<0.000000e+00> : tensor<1x1x1x1xf32>
%322 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%323 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%324 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%325 = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
%326 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%327 = func.call @fn_val_257(%318, %319, %320, %321, %322, %323, %324, %325, %326) : (tensor<1x4xi64>, tensor<8x4xf32>, tensor<1x1x4x4xf32>, tensor<1x1x1x1xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x8xf32>) -> tensor<1x4x8xf32>
return %327 : tensor<1x4x8xf32>
}
}