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