kumi 0.0.18 → 0.0.20

data/golden/multirank_hoisting/expected/ast.txt

@@ -0,0 +1,72 @@
+(Root
+  inputs: [
+    (InputDeclaration :batch :array
+      [
+        (InputDeclaration :b :hash
+          [
+            (InputDeclaration :mean :float)
+            (InputDeclaration :row :array
+              [
+                (InputDeclaration :r :hash
+                  [
+                    (InputDeclaration :scale :float)
+                    (InputDeclaration :col :array
+                      [
+                        (InputDeclaration :c :hash
+                          [
+                            (InputDeclaration :val :float)
+                          ]
+                        )
+                      ]
+                    )
+                  ]
+                )
+              ]
+            )
+          ]
+        )
+      ]
+    )
+    (InputDeclaration :global_offset :float)
+  ]
+  values: [
+    (ValueDeclaration :global_offset_plus
+      (CallExpression :add
+        (InputReference :global_offset)
+        (Literal 1.0)
+      )
+    )
+    (ValueDeclaration :batch_bias
+      (CallExpression :add
+        (InputElementReference batch.b.mean)
+        (DeclarationReference :global_offset_plus)
+      )
+    )
+    (ValueDeclaration :row_scale2
+      (CallExpression :multiply
+        (InputElementReference batch.b.row.r.scale)
+        (Literal 2.0)
+      )
+    )
+    (ValueDeclaration :elem_affine
+      (CallExpression :add
+        (CallExpression :multiply
+          (InputElementReference batch.b.row.r.col.c.val)
+          (DeclarationReference :row_scale2)
+        )
+        (DeclarationReference :batch_bias)
+      )
+    )
+    (ValueDeclaration :row_sum_affine
+      (CallExpression :sum
+        (DeclarationReference :elem_affine)
+      )
+    )
+    (ValueDeclaration :batch_total_affine
+      (CallExpression :sum
+        (DeclarationReference :row_sum_affine)
+      )
+    )
+  ]
+  traits: []
+)

data/golden/multirank_hoisting/expected/input_plan.txt

@@ -0,0 +1,18 @@
+batch: array (array) access_mode=element
+  └─> b: element_access
+  b: hash (hash)
+    └─> mean: property_access
+    └─> row: property_access → array_loop
+    mean: float
+    row: array (array) access_mode=element
+      └─> r: element_access
+      r: hash (hash)
+        └─> scale: property_access
+        └─> col: property_access → array_loop
+        scale: float
+        col: array (array) access_mode=element
+          └─> c: element_access
+          c: hash (hash)
+            └─> val: property_access
+            val: float
+global_offset: float

data/golden/multirank_hoisting/expected/lir_00_unoptimized.txt

@@ -0,0 +1,75 @@
+(LIR
+  (Declaration global_offset_plus
+    %t1 = load_input "global_offset" :: float
+    %t2 = const 1.0 :: float
+    %t3 = call core.add(%t1, %t2) :: float
+    yield %t3
+  )
+  (Declaration batch_bias
+    %t4 = load_input "batch" :: array
+    loop batch id=L1 in %t4 as el=%batch_el_5, idx=%batch_i_6
+      %t7 = load_field batch_el_5["mean"] :: float
+      %t8 = load_decl global_offset_plus :: float
+      %t9 = call core.add(%t7, %t8) :: float
+      yield %t9
+    end_loop
+  )
+  (Declaration row_scale2
+    %t10 = load_input "batch" :: array
+    loop batch id=L2 in %t10 as el=%batch_el_11, idx=%batch_i_12
+      %t13 = load_field batch_el_11["row"] :: any
+      loop row id=L3 in %t13 as el=%row_el_14, idx=%row_i_15
+        %t16 = load_field row_el_14["scale"] :: float
+        %t17 = const 2.0 :: float
+        %t18 = call core.mul(%t16, %t17) :: float
+        yield %t18
+      end_loop
+    end_loop
+  )
+  (Declaration elem_affine
+    %t19 = load_input "batch" :: array
+    loop batch id=L4 in %t19 as el=%batch_el_20, idx=%batch_i_21
+      %t22 = load_field batch_el_20["row"] :: any
+      loop row id=L5 in %t22 as el=%row_el_23, idx=%row_i_24
+        %t25 = load_field row_el_23["col"] :: any
+        loop col id=L6 in %t25 as el=%col_el_26, idx=%col_i_27
+          %t28 = load_field col_el_26["val"] :: float
+          %t29 = load_decl row_scale2 axes=[batch, row] :: float
+          %t30 = call core.mul(%t28, %t29) :: float
+          %t31 = load_decl batch_bias axes=[batch] :: float
+          %t32 = call core.add(%t30, %t31) :: float
+          yield %t32
+        end_loop
+      end_loop
+    end_loop
+  )
+  (Declaration row_sum_affine
+    %t33 = load_input "batch" :: array
+    loop batch id=L7 in %t33 as el=%batch_el_34, idx=%batch_i_35
+      %t36 = load_field batch_el_34["row"] :: any
+      loop row id=L8 in %t36 as el=%row_el_37, idx=%row_i_38
+        %acc_39 = decl_acc :: float
+        %t40 = load_field row_el_37["col"] :: any
+        loop col id=L9 in %t40 as el=%col_el_41, idx=%col_i_42
+          %t43 = load_decl elem_affine axes=[batch, row, col] :: float
+          %acc_39 = acc_add agg.sum(%acc_39, %t43) :: float
+        end_loop
+        %t44 = acc_load %acc_39 :: float
+        yield %t44
+      end_loop
+    end_loop
+  )
+  (Declaration batch_total_affine
+    %t45 = load_input "batch" :: array
+    loop batch id=L10 in %t45 as el=%batch_el_46, idx=%batch_i_47
+      %acc_48 = decl_acc :: float
+      %t49 = load_field batch_el_46["row"] :: any
+      loop row id=L11 in %t49 as el=%row_el_50, idx=%row_i_51
+        %t52 = load_decl row_sum_affine axes=[batch, row] :: float
+        %acc_48 = acc_add agg.sum(%acc_48, %t52) :: float
+      end_loop
+      %t53 = acc_load %acc_48 :: float
+      yield %t53
+    end_loop
+  )
+)

data/golden/multirank_hoisting/expected/lir_01_hoist_scalar_references.txt

@@ -0,0 +1,75 @@
+(LIR
+  (Declaration global_offset_plus
+    %t1 = load_input "global_offset" :: float
+    %t2 = const 1.0 :: float
+    %t3 = call core.add(%t1, %t2) :: float
+    yield %t3
+  )
+  (Declaration batch_bias
+    %t4 = load_input "batch" :: array
+    %t8 = load_decl global_offset_plus :: float
+    loop batch id=L1 in %t4 as el=%batch_el_5, idx=%batch_i_6
+      %t7 = load_field batch_el_5["mean"] :: float
+      %t9 = call core.add(%t7, %t8) :: float
+      yield %t9
+    end_loop
+  )
+  (Declaration row_scale2
+    %t10 = load_input "batch" :: array
+    loop batch id=L2 in %t10 as el=%batch_el_11, idx=%batch_i_12
+      %t13 = load_field batch_el_11["row"] :: any
+      loop row id=L3 in %t13 as el=%row_el_14, idx=%row_i_15
+        %t16 = load_field row_el_14["scale"] :: float
+        %t17 = const 2.0 :: float
+        %t18 = call core.mul(%t16, %t17) :: float
+        yield %t18
+      end_loop
+    end_loop
+  )
+  (Declaration elem_affine
+    %t19 = load_input "batch" :: array
+    loop batch id=L4 in %t19 as el=%batch_el_20, idx=%batch_i_21
+      %t22 = load_field batch_el_20["row"] :: any
+      loop row id=L5 in %t22 as el=%row_el_23, idx=%row_i_24
+        %t25 = load_field row_el_23["col"] :: any
+        loop col id=L6 in %t25 as el=%col_el_26, idx=%col_i_27
+          %t28 = load_field col_el_26["val"] :: float
+          %t29 = load_decl row_scale2 axes=[batch, row] :: float
+          %t30 = call core.mul(%t28, %t29) :: float
+          %t31 = load_decl batch_bias axes=[batch] :: float
+          %t32 = call core.add(%t30, %t31) :: float
+          yield %t32
+        end_loop
+      end_loop
+    end_loop
+  )
+  (Declaration row_sum_affine
+    %t33 = load_input "batch" :: array
+    loop batch id=L7 in %t33 as el=%batch_el_34, idx=%batch_i_35
+      %t36 = load_field batch_el_34["row"] :: any
+      loop row id=L8 in %t36 as el=%row_el_37, idx=%row_i_38
+        %acc_39 = decl_acc :: float
+        %t40 = load_field row_el_37["col"] :: any
+        loop col id=L9 in %t40 as el=%col_el_41, idx=%col_i_42
+          %t43 = load_decl elem_affine axes=[batch, row, col] :: float
+          %acc_39 = acc_add agg.sum(%acc_39, %t43) :: float
+        end_loop
+        %t44 = acc_load %acc_39 :: float
+        yield %t44
+      end_loop
+    end_loop
+  )
+  (Declaration batch_total_affine
+    %t45 = load_input "batch" :: array
+    loop batch id=L10 in %t45 as el=%batch_el_46, idx=%batch_i_47
+      %acc_48 = decl_acc :: float
+      %t49 = load_field batch_el_46["row"] :: any
+      loop row id=L11 in %t49 as el=%row_el_50, idx=%row_i_51
+        %t52 = load_decl row_sum_affine axes=[batch, row] :: float
+        %acc_48 = acc_add agg.sum(%acc_48, %t52) :: float
+      end_loop
+      %t53 = acc_load %acc_48 :: float
+      yield %t53
+    end_loop
+  )
+)

data/golden/multirank_hoisting/expected/lir_02_inlined.txt

@@ -0,0 +1,126 @@
+(LIR
+  (Declaration global_offset_plus
+    %t1 = load_input "global_offset" :: float
+    %t2 = const 1.0 :: float
+    %t3 = call core.add(%t1, %t2) :: float
+    yield %t3
+  )
+  (Declaration batch_bias
+    %t4 = load_input "batch" :: array
+    %t54 = load_input "global_offset" :: float
+    %t55 = const 1.0 :: float
+    %t56 = call core.add(%t54, %t55) :: float
+    loop batch id=L1 in %t4 as el=%batch_el_5, idx=%batch_i_6
+      %t7 = load_field batch_el_5["mean"] :: float
+      %t9 = call core.add(%t7, %t56) :: float
+      yield %t9
+    end_loop
+  )
+  (Declaration row_scale2
+    %t10 = load_input "batch" :: array
+    loop batch id=L2 in %t10 as el=%batch_el_11, idx=%batch_i_12
+      %t13 = load_field batch_el_11["row"] :: any
+      loop row id=L3 in %t13 as el=%row_el_14, idx=%row_i_15
+        %t16 = load_field row_el_14["scale"] :: float
+        %t17 = const 2.0 :: float
+        %t18 = call core.mul(%t16, %t17) :: float
+        yield %t18
+      end_loop
+    end_loop
+  )
+  (Declaration elem_affine
+    %t19 = load_input "batch" :: array
+    loop batch id=L4 in %t19 as el=%batch_el_20, idx=%batch_i_21
+      %t22 = load_field batch_el_20["row"] :: any
+      %t81 = load_input "global_offset" :: float
+      %t82 = const 1.0 :: float
+      %t83 = call core.add(%t81, %t82) :: float
+      loop row id=L5 in %t22 as el=%row_el_23, idx=%row_i_24
+        %t25 = load_field row_el_23["col"] :: any
+        %t57 = load_input "batch" :: array
+        %t58 = load_field batch_el_20["row"] :: any
+        %t59 = load_field row_el_23["scale"] :: float
+        %t60 = const 2.0 :: float
+        %t61 = call core.mul(%t59, %t60) :: float
+        %t62 = load_input "batch" :: array
+        %t64 = load_field batch_el_20["mean"] :: float
+        %t65 = call core.add(%t64, %t83) :: float
+        loop col id=L6 in %t25 as el=%col_el_26, idx=%col_i_27
+          %t28 = load_field col_el_26["val"] :: float
+          %t30 = call core.mul(%t28, %t61) :: float
+          %t32 = call core.add(%t30, %t65) :: float
+          yield %t32
+        end_loop
+      end_loop
+    end_loop
+  )
+  (Declaration row_sum_affine
+    %t33 = load_input "batch" :: array
+    loop batch id=L7 in %t33 as el=%batch_el_34, idx=%batch_i_35
+      %t36 = load_field batch_el_34["row"] :: any
+      loop row id=L8 in %t36 as el=%row_el_37, idx=%row_i_38
+        %acc_39 = decl_acc :: float
+        %t40 = load_field row_el_37["col"] :: any
+        %t84 = load_input "batch" :: array
+        %t85 = load_field batch_el_34["row"] :: any
+        %t86 = load_field row_el_37["scale"] :: float
+        %t87 = const 2.0 :: float
+        %t88 = call core.mul(%t86, %t87) :: float
+        %t89 = load_input "batch" :: array
+        %t90 = load_input "global_offset" :: float
+        %t91 = const 1.0 :: float
+        %t92 = call core.add(%t90, %t91) :: float
+        %t93 = load_field batch_el_34["mean"] :: float
+        %t94 = call core.add(%t93, %t92) :: float
+        loop col id=L9 in %t40 as el=%col_el_41, idx=%col_i_42
+          %t66 = load_input "batch" :: array
+          %t67 = load_field batch_el_34["row"] :: any
+          %t68 = load_field row_el_37["col"] :: any
+          %t69 = load_field col_el_41["val"] :: float
+          %t71 = call core.mul(%t69, %t88) :: float
+          %t73 = call core.add(%t71, %t94) :: float
+          %acc_39 = acc_add agg.sum(%acc_39, %t73) :: float
+        end_loop
+        %t44 = acc_load %acc_39 :: float
+        yield %t44
+      end_loop
+    end_loop
+  )
+  (Declaration batch_total_affine
+    %t45 = load_input "batch" :: array
+    loop batch id=L10 in %t45 as el=%batch_el_46, idx=%batch_i_47
+      %acc_48 = decl_acc :: float
+      %t49 = load_field batch_el_46["row"] :: any
+      loop row id=L11 in %t49 as el=%row_el_50, idx=%row_i_51
+        %t74 = load_input "batch" :: array
+        %t75 = load_field batch_el_46["row"] :: any
+        %acc77 = decl_acc :: float
+        %t78 = load_field row_el_50["col"] :: any
+        %t110 = load_input "global_offset" :: float
+        %t111 = const 1.0 :: float
+        %t112 = call core.add(%t110, %t111) :: float
+        loop col id=L12 in %t78 as el=%col_el_41, idx=%col_i_42
+          %t95 = load_input "batch" :: array
+          %t96 = load_field batch_el_46["row"] :: any
+          %t97 = load_field row_el_50["col"] :: any
+          %t98 = load_input "batch" :: array
+          %t99 = load_field batch_el_46["row"] :: any
+          %t100 = load_field row_el_50["scale"] :: float
+          %t101 = const 2.0 :: float
+          %t102 = call core.mul(%t100, %t101) :: float
+          %t103 = load_input "batch" :: array
+          %t105 = load_field batch_el_46["mean"] :: float
+          %t106 = call core.add(%t105, %t112) :: float
+          %t107 = load_field col_el_41["val"] :: float
+          %t108 = call core.mul(%t107, %t102) :: float
+          %t109 = call core.add(%t108, %t106) :: float
+          %acc77 = acc_add agg.sum(%acc77, %t109) :: float
+        end_loop
+        %t80 = acc_load %acc77 :: float
+        %acc_48 = acc_add agg.sum(%acc_48, %t80) :: float
+      end_loop
+      %t53 = acc_load %acc_48 :: float
+      yield %t53
+    end_loop
+  )
+)

data/golden/multirank_hoisting/expected/lir_03_cse.txt

@@ -0,0 +1,109 @@
+(LIR
+  (Declaration global_offset_plus
+    %t1 = load_input "global_offset" :: float
+    %t2 = const 1.0 :: float
+    %t3 = call core.add(%t1, %t2) :: float
+    yield %t3
+  )
+  (Declaration batch_bias
+    %t4 = load_input "batch" :: array
+    %t54 = load_input "global_offset" :: float
+    %t55 = const 1.0 :: float
+    %t56 = call core.add(%t54, %t55) :: float
+    loop batch id=L1 in %t4 as el=%batch_el_5, idx=%batch_i_6
+      %t7 = load_field batch_el_5["mean"] :: float
+      %t9 = call core.add(%t7, %t56) :: float
+      yield %t9
+    end_loop
+  )
+  (Declaration row_scale2
+    %t10 = load_input "batch" :: array
+    loop batch id=L2 in %t10 as el=%batch_el_11, idx=%batch_i_12
+      %t13 = load_field batch_el_11["row"] :: any
+      loop row id=L3 in %t13 as el=%row_el_14, idx=%row_i_15
+        %t16 = load_field row_el_14["scale"] :: float
+        %t17 = const 2.0 :: float
+        %t18 = call core.mul(%t16, %t17) :: float
+        yield %t18
+      end_loop
+    end_loop
+  )
+  (Declaration elem_affine
+    %t19 = load_input "batch" :: array
+    loop batch id=L4 in %t19 as el=%batch_el_20, idx=%batch_i_21
+      %t22 = load_field batch_el_20["row"] :: any
+      %t81 = load_input "global_offset" :: float
+      %t82 = const 1.0 :: float
+      %t83 = call core.add(%t81, %t82) :: float
+      loop row id=L5 in %t22 as el=%row_el_23, idx=%row_i_24
+        %t25 = load_field row_el_23["col"] :: any
+        %t59 = load_field row_el_23["scale"] :: float
+        %t60 = const 2.0 :: float
+        %t61 = call core.mul(%t59, %t60) :: float
+        %t64 = load_field batch_el_20["mean"] :: float
+        %t65 = call core.add(%t64, %t83) :: float
+        loop col id=L6 in %t25 as el=%col_el_26, idx=%col_i_27
+          %t28 = load_field col_el_26["val"] :: float
+          %t30 = call core.mul(%t28, %t61) :: float
+          %t32 = call core.add(%t30, %t65) :: float
+          yield %t32
+        end_loop
+      end_loop
+    end_loop
+  )
+  (Declaration row_sum_affine
+    %t33 = load_input "batch" :: array
+    loop batch id=L7 in %t33 as el=%batch_el_34, idx=%batch_i_35
+      %t36 = load_field batch_el_34["row"] :: any
+      loop row id=L8 in %t36 as el=%row_el_37, idx=%row_i_38
+        %acc_39 = decl_acc :: float
+        %t40 = load_field row_el_37["col"] :: any
+        %t86 = load_field row_el_37["scale"] :: float
+        %t87 = const 2.0 :: float
+        %t88 = call core.mul(%t86, %t87) :: float
+        %t90 = load_input "global_offset" :: float
+        %t91 = const 1.0 :: float
+        %t92 = call core.add(%t90, %t91) :: float
+        %t93 = load_field batch_el_34["mean"] :: float
+        %t94 = call core.add(%t93, %t92) :: float
+        loop col id=L9 in %t40 as el=%col_el_41, idx=%col_i_42
+          %t69 = load_field col_el_41["val"] :: float
+          %t71 = call core.mul(%t69, %t88) :: float
+          %t73 = call core.add(%t71, %t94) :: float
+          %acc_39 = acc_add agg.sum(%acc_39, %t73) :: float
+        end_loop
+        %t44 = acc_load %acc_39 :: float
+        yield %t44
+      end_loop
+    end_loop
+  )
+  (Declaration batch_total_affine
+    %t45 = load_input "batch" :: array
+    loop batch id=L10 in %t45 as el=%batch_el_46, idx=%batch_i_47
+      %acc_48 = decl_acc :: float
+      %t49 = load_field batch_el_46["row"] :: any
+      loop row id=L11 in %t49 as el=%row_el_50, idx=%row_i_51
+        %acc77 = decl_acc :: float
+        %t78 = load_field row_el_50["col"] :: any
+        %t110 = load_input "global_offset" :: float
+        %t111 = const 1.0 :: float
+        %t112 = call core.add(%t110, %t111) :: float
+        loop col id=L12 in %t78 as el=%col_el_41, idx=%col_i_42
+          %t100 = load_field row_el_50["scale"] :: float
+          %t101 = const 2.0 :: float
+          %t102 = call core.mul(%t100, %t101) :: float
+          %t105 = load_field batch_el_46["mean"] :: float
+          %t106 = call core.add(%t105, %t112) :: float
+          %t107 = load_field col_el_41["val"] :: float
+          %t108 = call core.mul(%t107, %t102) :: float
+          %t109 = call core.add(%t108, %t106) :: float
+          %acc77 = acc_add agg.sum(%acc77, %t109) :: float
+        end_loop
+        %t80 = acc_load %acc77 :: float
+        %acc_48 = acc_add agg.sum(%acc_48, %t80) :: float
+      end_loop
+      %t53 = acc_load %acc_48 :: float
+      yield %t53
+    end_loop
+  )
+)

data/golden/multirank_hoisting/expected/lir_04_1_loop_fusion.txt

@@ -0,0 +1,109 @@
+(LIR
+  (Declaration global_offset_plus
+    %t1 = load_input "global_offset" :: float
+    %t2 = const 1.0 :: float
+    %t3 = call core.add(%t1, %t2) :: float
+    yield %t3
+  )
+  (Declaration batch_bias
+    %t4 = load_input "batch" :: array
+    %t54 = load_input "global_offset" :: float
+    %t55 = const 1.0 :: float
+    %t56 = call core.add(%t54, %t55) :: float
+    loop batch id=L1 in %t4 as el=%batch_el_5, idx=%batch_i_6
+      %t7 = load_field batch_el_5["mean"] :: float
+      %t9 = call core.add(%t7, %t56) :: float
+      yield %t9
+    end_loop
+  )
+  (Declaration row_scale2
+    %t10 = load_input "batch" :: array
+    loop batch id=L2 in %t10 as el=%batch_el_11, idx=%batch_i_12
+      %t13 = load_field batch_el_11["row"] :: any
+      loop row id=L3 in %t13 as el=%row_el_14, idx=%row_i_15
+        %t16 = load_field row_el_14["scale"] :: float
+        %t17 = const 2.0 :: float
+        %t18 = call core.mul(%t16, %t17) :: float
+        yield %t18
+      end_loop
+    end_loop
+  )
+  (Declaration elem_affine
+    %t19 = load_input "batch" :: array
+    loop batch id=L4 in %t19 as el=%batch_el_20, idx=%batch_i_21
+      %t22 = load_field batch_el_20["row"] :: any
+      %t81 = load_input "global_offset" :: float
+      %t82 = const 1.0 :: float
+      %t83 = call core.add(%t81, %t82) :: float
+      loop row id=L5 in %t22 as el=%row_el_23, idx=%row_i_24
+        %t25 = load_field row_el_23["col"] :: any
+        %t59 = load_field row_el_23["scale"] :: float
+        %t60 = const 2.0 :: float
+        %t61 = call core.mul(%t59, %t60) :: float
+        %t64 = load_field batch_el_20["mean"] :: float
+        %t65 = call core.add(%t64, %t83) :: float
+        loop col id=L6 in %t25 as el=%col_el_26, idx=%col_i_27
+          %t28 = load_field col_el_26["val"] :: float
+          %t30 = call core.mul(%t28, %t61) :: float
+          %t32 = call core.add(%t30, %t65) :: float
+          yield %t32
+        end_loop
+      end_loop
+    end_loop
+  )
+  (Declaration row_sum_affine
+    %t33 = load_input "batch" :: array
+    loop batch id=L7 in %t33 as el=%batch_el_34, idx=%batch_i_35
+      %t36 = load_field batch_el_34["row"] :: any
+      loop row id=L8 in %t36 as el=%row_el_37, idx=%row_i_38
+        %acc_39 = decl_acc :: float
+        %t40 = load_field row_el_37["col"] :: any
+        %t86 = load_field row_el_37["scale"] :: float
+        %t87 = const 2.0 :: float
+        %t88 = call core.mul(%t86, %t87) :: float
+        %t90 = load_input "global_offset" :: float
+        %t91 = const 1.0 :: float
+        %t92 = call core.add(%t90, %t91) :: float
+        %t93 = load_field batch_el_34["mean"] :: float
+        %t94 = call core.add(%t93, %t92) :: float
+        loop col id=L9 in %t40 as el=%col_el_41, idx=%col_i_42
+          %t69 = load_field col_el_41["val"] :: float
+          %t71 = call core.mul(%t69, %t88) :: float
+          %t73 = call core.add(%t71, %t94) :: float
+          %acc_39 = acc_add agg.sum(%acc_39, %t73) :: float
+        end_loop
+        %t44 = acc_load %acc_39 :: float
+        yield %t44
+      end_loop
+    end_loop
+  )
+  (Declaration batch_total_affine
+    %t45 = load_input "batch" :: array
+    loop batch id=L10 in %t45 as el=%batch_el_46, idx=%batch_i_47
+      %acc_48 = decl_acc :: float
+      %t49 = load_field batch_el_46["row"] :: any
+      loop row id=L11 in %t49 as el=%row_el_50, idx=%row_i_51
+        %acc77 = decl_acc :: float
+        %t78 = load_field row_el_50["col"] :: any
+        %t110 = load_input "global_offset" :: float
+        %t111 = const 1.0 :: float
+        %t112 = call core.add(%t110, %t111) :: float
+        loop col id=L12 in %t78 as el=%col_el_41, idx=%col_i_42
+          %t100 = load_field row_el_50["scale"] :: float
+          %t101 = const 2.0 :: float
+          %t102 = call core.mul(%t100, %t101) :: float
+          %t105 = load_field batch_el_46["mean"] :: float
+          %t106 = call core.add(%t105, %t112) :: float
+          %t107 = load_field col_el_41["val"] :: float
+          %t108 = call core.mul(%t107, %t102) :: float
+          %t109 = call core.add(%t108, %t106) :: float
+          %acc77 = acc_add agg.sum(%acc77, %t109) :: float
+        end_loop
+        %t80 = acc_load %acc77 :: float
+        %acc_48 = acc_add agg.sum(%acc_48, %t80) :: float
+      end_loop
+      %t53 = acc_load %acc_48 :: float
+      yield %t53
+    end_loop
+  )
+)