Example
adamw_step.kina
let adamw_step = fn(
w: tensor<F32, [4, 8], @host>,
g: tensor<F32, [4, 8], @host>,
m: tensor<F32, [4, 8], @host>,
v: tensor<F32, [4, 8], @host>,
lr: tensor<F32, [1, 1], @host>,
beta1: tensor<F32, [1, 1], @host>,
beta2: tensor<F32, [1, 1], @host>,
eps: tensor<F32, [1, 1], @host>,
wd: tensor<F32, [1, 1], @host>,
beta1_pow: tensor<F32, [1, 1], @host>,
beta2_pow: tensor<F32, [1, 1], @host>,
one: tensor<F32, [1, 1], @host>
) -> (
tensor<F32, [4, 8], @host>,
tensor<F32, [4, 8], @host>,
tensor<F32, [4, 8], @host>
) {
let one_minus_beta1 = one - beta1 in
let one_minus_beta2 = one - beta2 in
let beta1_b = beta1.broadcast([4, 8]) in
let beta2_b = beta2.broadcast([4, 8]) in
let one_minus_beta1_b = one_minus_beta1.broadcast([4, 8]) in
let one_minus_beta2_b = one_minus_beta2.broadcast([4, 8]) in
let m_keep = beta1_b * m in
let m_grad = one_minus_beta1_b * g in
let m_next = m_keep + m_grad in
let g2 = g * g in
let v_keep = beta2_b * v in
let v_grad = one_minus_beta2_b * g2 in
let v_next = v_keep + v_grad in
let bc1 = one - beta1_pow in
let bc2 = one - beta2_pow in
let v_sqrt = sqrt(v_next) in
let bc2_sqrt = sqrt(bc2) in
let bc2_sqrt_b = bc2_sqrt.broadcast([4, 8]) in
let v_norm = v_sqrt / bc2_sqrt_b in
let eps_b = eps.broadcast([4, 8]) in
let denom = v_norm + eps_b in
let step_size = lr / bc1 in
let step_size_b = step_size.broadcast([4, 8]) in
let m_over_denom = m_next / denom in
let adam_step = step_size_b * m_over_denom in
let decay = one - lr * wd in
let decay_b = decay.broadcast([4, 8]) in
let w_decayed = w * decay_b in
let w_next = w_decayed - adam_step in
(w_next, m_next, v_next)
} in
let w = tensor<F32, [4, 8], @host> in
let g = tensor<F32, [4, 8], @host> in
let m = tensor<F32, [4, 8], @host> in
let v = tensor<F32, [4, 8], @host> in
let lr = tensor<F32, [1, 1], @host> in
let beta1 = tensor<F32, [1, 1], @host> in
let beta2 = tensor<F32, [1, 1], @host> in
let eps = tensor<F32, [1, 1], @host> in
let wd = tensor<F32, [1, 1], @host> in
let beta1_pow = tensor<F32, [1, 1], @host> in
let beta2_pow = tensor<F32, [1, 1], @host> in
let one = tensor<F32, [1, 1], @host> in
adamw_step(w, g, m, v, lr, beta1, beta2, eps, wd, beta1_pow, beta2_pow, one)
module {
func.func @fn_val_0(%1: tensor<1x32x64xf32>, %2: tensor<1x1x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%3 = tensor.empty() : tensor<1x32x64xf32>
%4 = linalg.mul ins(%1, %1 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%3 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%7 = arith.constant 0.000000e+00 : f32
%5 = tensor.empty() : tensor<1x32x1xf32>
%6 = linalg.fill ins(%7 : f32) outs(%5 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%8 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%4 : tensor<1x32x64xf32>) outs(%6 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%9 = arith.addf %a0, %a1 : f32
linalg.yield %9 : f32
} -> tensor<1x32x1xf32>
%10 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%11 = tensor.empty() : tensor<1x32x1xf32>
%12 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%10 : tensor<1x1x1xf32>) outs(%11 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%13 = tensor.empty() : tensor<1x32x1xf32>
%14 = linalg.mul ins(%8, %12 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%13 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%15 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%16 = tensor.empty() : tensor<1x32x1xf32>
%17 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%15 : tensor<1x1x1xf32>) outs(%16 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x1xf32>
%18 = tensor.empty() : tensor<1x32x1xf32>
%19 = linalg.add ins(%14, %17 : tensor<1x32x1xf32>, tensor<1x32x1xf32>) outs(%18 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%23 = tensor.empty() : tensor<1x32x1xf32>
%24 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%19 : tensor<1x32x1xf32>) outs(%23 : tensor<1x32x1xf32>) {
^bb0(%x: f32, %_: f32):
%20 = math.sqrt %x : f32
%21 = arith.constant 1.000000e+00 : f32
%22 = arith.divf %21, %20 : f32
linalg.yield %22 : f32
} -> tensor<1x32x1xf32>
%25 = tensor.empty() : tensor<1x32x64xf32>
%26 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%24 : tensor<1x32x1xf32>) outs(%25 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%27 = tensor.empty() : tensor<1x32x64xf32>
%28 = linalg.mul ins(%1, %26 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%27 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%29 = tensor.empty() : tensor<1x32x64xf32>
%30 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%2 : tensor<1x1x64xf32>) outs(%29 : tensor<1x32x64xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x64xf32>
%31 = tensor.empty() : tensor<1x32x64xf32>
%32 = linalg.mul ins(%28, %30 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%31 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %32 : tensor<1x32x64xf32>
}
func.func @fn_val_33(%34: tensor<1x32x64xf32>, %35: tensor<1x64x128xf32>, %36: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%37 = arith.constant dense<0.000000e+00> : tensor<1x32x128xf32>
%38 = linalg.batch_matmul ins(%34, %35 : tensor<1x32x64xf32>, tensor<1x64x128xf32>) outs(%37 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%44 = tensor.empty() : tensor<1x32x128xf32>
%45 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%38 : tensor<1x32x128xf32>) outs(%44 : tensor<1x32x128xf32>) {
^bb0(%a0: f32, %a1: f32):
%39 = arith.constant 1.000000e+00 : f32
%40 = arith.negf %a0 : f32
%41 = math.exp %40 : f32
%42 = arith.addf %39, %41 : f32
%43 = arith.divf %39, %42 : f32
linalg.yield %43 : f32
} -> tensor<1x32x128xf32>
%46 = tensor.empty() : tensor<1x32x128xf32>
%47 = linalg.mul ins(%38, %45 : tensor<1x32x128xf32>, tensor<1x32x128xf32>) outs(%46 : tensor<1x32x128xf32>) -> tensor<1x32x128xf32>
%48 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%49 = linalg.batch_matmul ins(%47, %36 : tensor<1x32x128xf32>, tensor<1x128x64xf32>) outs(%48 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %49 : tensor<1x32x64xf32>
}
func.func @fn_val_50(%51: tensor<1x32x64xf32>, %52: tensor<1x64x64xf32>, %53: tensor<1x64x64xf32>, %54: tensor<1x64x64xf32>, %55: tensor<1x64x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%56 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%57 = linalg.batch_matmul ins(%51, %52 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%56 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%58 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%59 = linalg.batch_matmul ins(%51, %53 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%58 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%60 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%61 = linalg.batch_matmul ins(%51, %54 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%60 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%62 = tensor.empty() : tensor<1x64x32xf32>
%63 = linalg.transpose ins(%59 : tensor<1x32x64xf32>) outs(%62 : tensor<1x64x32xf32>) permutation = [0, 2, 1]
%64 = arith.constant dense<0.000000e+00> : tensor<1x32x32xf32>
%65 = linalg.batch_matmul ins(%57, %63 : tensor<1x32x64xf32>, tensor<1x64x32xf32>) outs(%64 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%66 = arith.constant dense<0.000000e+00> : tensor<1x1x1xf32>
%67 = tensor.empty() : tensor<1x32x32xf32>
%68 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, 0, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%66 : tensor<1x1x1xf32>) outs(%67 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%69 = tensor.empty() : tensor<1x32x32xf32>
%70 = linalg.mul ins(%65, %68 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%69 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%73 = arith.constant 0xFF800000 : f32
%71 = tensor.empty() : tensor<1x32x1xf32>
%72 = linalg.fill ins(%73 : f32) outs(%71 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%74 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%70 : tensor<1x32x32xf32>) outs(%72 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%75 = arith.maximumf %a0, %a1 : f32
linalg.yield %75 : f32
} -> tensor<1x32x1xf32>
%76 = tensor.empty() : tensor<1x32x32xf32>
%77 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%74 : tensor<1x32x1xf32>) outs(%76 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%78 = tensor.empty() : tensor<1x32x32xf32>
%79 = linalg.sub ins(%70, %77 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%78 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%80 = tensor.empty() : tensor<1x32x32xf32>
%81 = linalg.exp ins(%79 : tensor<1x32x32xf32>) outs(%80 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%84 = arith.constant 0xFF800000 : f32
%82 = tensor.empty() : tensor<1x32x1xf32>
%83 = linalg.fill ins(%84 : f32) outs(%82 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%85 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%70 : tensor<1x32x32xf32>) outs(%83 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%86 = arith.maximumf %a0, %a1 : f32
linalg.yield %86 : f32
} -> tensor<1x32x1xf32>
%87 = tensor.empty() : tensor<1x32x32xf32>
%88 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%85 : tensor<1x32x1xf32>) outs(%87 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%89 = tensor.empty() : tensor<1x32x32xf32>
%90 = linalg.sub ins(%70, %88 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%89 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%91 = tensor.empty() : tensor<1x32x32xf32>
%92 = linalg.exp ins(%90 : tensor<1x32x32xf32>) outs(%91 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%95 = arith.constant 0.000000e+00 : f32
%93 = tensor.empty() : tensor<1x32x1xf32>
%94 = linalg.fill ins(%95 : f32) outs(%93 : tensor<1x32x1xf32>) -> tensor<1x32x1xf32>
%96 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 0)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%92 : tensor<1x32x32xf32>) outs(%94 : tensor<1x32x1xf32>) {
^bb0(%a0: f32, %a1: f32):
%97 = arith.addf %a0, %a1 : f32
linalg.yield %97 : f32
} -> tensor<1x32x1xf32>
%98 = tensor.empty() : tensor<1x32x32xf32>
%99 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (0, d1, 0)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]
} ins(%96 : tensor<1x32x1xf32>) outs(%98 : tensor<1x32x32xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<1x32x32xf32>
%100 = tensor.empty() : tensor<1x32x32xf32>
%101 = linalg.div ins(%81, %99 : tensor<1x32x32xf32>, tensor<1x32x32xf32>) outs(%100 : tensor<1x32x32xf32>) -> tensor<1x32x32xf32>
%102 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%103 = linalg.batch_matmul ins(%101, %61 : tensor<1x32x32xf32>, tensor<1x32x64xf32>) outs(%102 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%104 = arith.constant dense<0.000000e+00> : tensor<1x32x64xf32>
%105 = linalg.batch_matmul ins(%103, %55 : tensor<1x32x64xf32>, tensor<1x64x64xf32>) outs(%104 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %105 : tensor<1x32x64xf32>
}
func.func @fn_val_106(%107: tensor<1x32x64xf32>, %108: tensor<1x1x64xf32>, %109: tensor<1x64x64xf32>, %110: tensor<1x64x64xf32>, %111: tensor<1x64x64xf32>, %112: tensor<1x64x64xf32>, %113: tensor<1x1x64xf32>, %114: tensor<1x64x128xf32>, %115: tensor<1x128x64xf32>) -> tensor<1x32x64xf32> attributes {llvm.emit_c_interface} {
%116 = func.call @fn_val_0(%107, %108) : (tensor<1x32x64xf32>, tensor<1x1x64xf32>) -> tensor<1x32x64xf32>
%117 = func.call @fn_val_50(%116, %109, %110, %111, %112) : (tensor<1x32x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>, tensor<1x64x64xf32>) -> tensor<1x32x64xf32>
%118 = tensor.empty() : tensor<1x32x64xf32>
%119 = linalg.add ins(%107, %117 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%118 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
%120 = func.call @fn_val_0(%119, %113) : (tensor<1x32x64xf32>, tensor<1x1x64xf32>) -> tensor<1x32x64xf32>
%121 = func.call @fn_val_33(%120, %114, %115) : (tensor<1x32x64xf32>, tensor<1x64x128xf32>, tensor<1x128x64xf32>) -> tensor<1x32x64xf32>
%122 = tensor.empty() : tensor<1x32x64xf32>
%123 = linalg.add ins(%119, %121 : tensor<1x32x64xf32>, tensor<1x32x64xf32>) outs(%122 : tensor<1x32x64xf32>) -> tensor<1x32x64xf32>
return %123 : tensor<1x32x64xf32>
}
func.func @fn_val_124(%125: tensor<4x8xf32>, %126: tensor<4x8xf32>, %127: tensor<4x8xf32>, %128: tensor<4x8xf32>, %129: tensor<1x1xf32>, %130: tensor<1x1xf32>, %131: tensor<1x1xf32>, %132: tensor<1x1xf32>, %133: tensor<1x1xf32>, %134: tensor<1x1xf32>, %135: tensor<1x1xf32>, %136: tensor<1x1xf32>) -> (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>) attributes {llvm.emit_c_interface} {
%137 = tensor.empty() : tensor<1x1xf32>
%138 = linalg.sub ins(%136, %130 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%137 : tensor<1x1xf32>) -> tensor<1x1xf32>
%139 = tensor.empty() : tensor<1x1xf32>
%140 = linalg.sub ins(%136, %131 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%139 : tensor<1x1xf32>) -> tensor<1x1xf32>
%141 = tensor.empty() : tensor<4x8xf32>
%142 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%130 : tensor<1x1xf32>) outs(%141 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%143 = tensor.empty() : tensor<4x8xf32>
%144 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%131 : tensor<1x1xf32>) outs(%143 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%145 = tensor.empty() : tensor<4x8xf32>
%146 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%138 : tensor<1x1xf32>) outs(%145 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%147 = tensor.empty() : tensor<4x8xf32>
%148 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%140 : tensor<1x1xf32>) outs(%147 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%149 = tensor.empty() : tensor<4x8xf32>
%150 = linalg.mul ins(%142, %127 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%149 : tensor<4x8xf32>) -> tensor<4x8xf32>
%151 = tensor.empty() : tensor<4x8xf32>
%152 = linalg.mul ins(%146, %126 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%151 : tensor<4x8xf32>) -> tensor<4x8xf32>
%153 = tensor.empty() : tensor<4x8xf32>
%154 = linalg.add ins(%150, %152 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%153 : tensor<4x8xf32>) -> tensor<4x8xf32>
%155 = tensor.empty() : tensor<4x8xf32>
%156 = linalg.mul ins(%126, %126 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%155 : tensor<4x8xf32>) -> tensor<4x8xf32>
%157 = tensor.empty() : tensor<4x8xf32>
%158 = linalg.mul ins(%144, %128 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%157 : tensor<4x8xf32>) -> tensor<4x8xf32>
%159 = tensor.empty() : tensor<4x8xf32>
%160 = linalg.mul ins(%148, %156 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%159 : tensor<4x8xf32>) -> tensor<4x8xf32>
%161 = tensor.empty() : tensor<4x8xf32>
%162 = linalg.add ins(%158, %160 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%161 : tensor<4x8xf32>) -> tensor<4x8xf32>
%163 = tensor.empty() : tensor<1x1xf32>
%164 = linalg.sub ins(%136, %134 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%163 : tensor<1x1xf32>) -> tensor<1x1xf32>
%165 = tensor.empty() : tensor<1x1xf32>
%166 = linalg.sub ins(%136, %135 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%165 : tensor<1x1xf32>) -> tensor<1x1xf32>
%168 = tensor.empty() : tensor<4x8xf32>
%169 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%162 : tensor<4x8xf32>) outs(%168 : tensor<4x8xf32>) {
^bb0(%x: f32, %_: f32):
%167 = math.sqrt %x : f32
linalg.yield %167 : f32
} -> tensor<4x8xf32>
%171 = tensor.empty() : tensor<1x1xf32>
%172 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%166 : tensor<1x1xf32>) outs(%171 : tensor<1x1xf32>) {
^bb0(%x: f32, %_: f32):
%170 = math.sqrt %x : f32
linalg.yield %170 : f32
} -> tensor<1x1xf32>
%173 = tensor.empty() : tensor<4x8xf32>
%174 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%172 : tensor<1x1xf32>) outs(%173 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%175 = tensor.empty() : tensor<4x8xf32>
%176 = linalg.div ins(%169, %174 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%175 : tensor<4x8xf32>) -> tensor<4x8xf32>
%177 = tensor.empty() : tensor<4x8xf32>
%178 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%132 : tensor<1x1xf32>) outs(%177 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%179 = tensor.empty() : tensor<4x8xf32>
%180 = linalg.add ins(%176, %178 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%179 : tensor<4x8xf32>) -> tensor<4x8xf32>
%181 = tensor.empty() : tensor<1x1xf32>
%182 = linalg.div ins(%129, %164 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%181 : tensor<1x1xf32>) -> tensor<1x1xf32>
%183 = tensor.empty() : tensor<4x8xf32>
%184 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%182 : tensor<1x1xf32>) outs(%183 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%185 = tensor.empty() : tensor<4x8xf32>
%186 = linalg.div ins(%154, %180 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%185 : tensor<4x8xf32>) -> tensor<4x8xf32>
%187 = tensor.empty() : tensor<4x8xf32>
%188 = linalg.mul ins(%184, %186 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%187 : tensor<4x8xf32>) -> tensor<4x8xf32>
%189 = tensor.empty() : tensor<1x1xf32>
%190 = linalg.mul ins(%129, %133 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%189 : tensor<1x1xf32>) -> tensor<1x1xf32>
%191 = tensor.empty() : tensor<1x1xf32>
%192 = linalg.sub ins(%136, %190 : tensor<1x1xf32>, tensor<1x1xf32>) outs(%191 : tensor<1x1xf32>) -> tensor<1x1xf32>
%193 = tensor.empty() : tensor<4x8xf32>
%194 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (0, 0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]
} ins(%192 : tensor<1x1xf32>) outs(%193 : tensor<4x8xf32>) {
^bb0(%a0: f32, %a1: f32):
linalg.yield %a0 : f32
} -> tensor<4x8xf32>
%195 = tensor.empty() : tensor<4x8xf32>
%196 = linalg.mul ins(%125, %194 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%195 : tensor<4x8xf32>) -> tensor<4x8xf32>
%197 = tensor.empty() : tensor<4x8xf32>
%198 = linalg.sub ins(%196, %188 : tensor<4x8xf32>, tensor<4x8xf32>) outs(%197 : tensor<4x8xf32>) -> tensor<4x8xf32>
return %198, %154, %162 : tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>
}
func.func @main() -> (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>) {
%199 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%200 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%201 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%202 = arith.constant dense<0.000000e+00> : tensor<4x8xf32>
%203 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%204 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%205 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%206 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%207 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%208 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%209 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%210 = arith.constant dense<0.000000e+00> : tensor<1x1xf32>
%211, %212, %213 = func.call @fn_val_124(%199, %200, %201, %202, %203, %204, %205, %206, %207, %208, %209, %210) : (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>)
return %211, %212, %213 : tensor<4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>
}
}