From a0eb92e88b765895095e0db3a9364841713367e8 Mon Sep 17 00:00:00 2001
From: Jacob Lifshay <programmerjake@gmail.com>
Date: Mon, 31 Oct 2022 17:45:26 -0700
Subject: [PATCH] add pre_ra_insert_copies

---
 .../_tests/test_compiler_ir2.py               | 95 +++++++++++++++++++
 src/bigint_presentation_code/compiler_ir2.py  | 46 +++++++++
 2 files changed, 141 insertions(+)

diff --git a/src/bigint_presentation_code/_tests/test_compiler_ir2.py b/src/bigint_presentation_code/_tests/test_compiler_ir2.py
index dfa43e6..288aea6 100644
--- a/src/bigint_presentation_code/_tests/test_compiler_ir2.py
+++ b/src/bigint_presentation_code/_tests/test_compiler_ir2.py
@@ -66,6 +66,101 @@ class TestCompilerIR(unittest.TestCase):
             "outputs=(), name='st')",
         ])
 
+    def test_pre_ra_insert_copies(self):
+        fn, _arg = self.make_add_fn()
+        fn.pre_ra_insert_copies()
+        self.assertEqual([repr(i) for i in fn.ops], [
+            "Op(kind=OpKind.FuncArgR3, "
+            "inputs=[], "
+            "immediates=[], "
+            "outputs=(<arg#0: <I64>>,), name='arg')",
+            "Op(kind=OpKind.CopyFromReg, "
+            "inputs=[<arg#0: <I64>>], "
+            "immediates=[], "
+            "outputs=(<2#0: <I64>>,), name='2')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<vl#0: <VL_MAXVL>>,), name='vl')",
+            "Op(kind=OpKind.CopyToReg, "
+            "inputs=[<2#0: <I64>>], "
+            "immediates=[], "
+            "outputs=(<3#0: <I64>>,), name='3')",
+            "Op(kind=OpKind.SvLd, "
+            "inputs=[<3#0: <I64>>, <vl#0: <VL_MAXVL>>], "
+            "immediates=[0], "
+            "outputs=(<ld#0: <I64*32>>,), name='ld')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<4#0: <VL_MAXVL>>,), name='4')",
+            "Op(kind=OpKind.VecCopyFromReg, "
+            "inputs=[<ld#0: <I64*32>>, <4#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<5#0: <I64*32>>,), name='5')",
+            "Op(kind=OpKind.SvLI, "
+            "inputs=[<vl#0: <VL_MAXVL>>], "
+            "immediates=[0], "
+            "outputs=(<li#0: <I64*32>>,), name='li')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<6#0: <VL_MAXVL>>,), name='6')",
+            "Op(kind=OpKind.VecCopyFromReg, "
+            "inputs=[<li#0: <I64*32>>, <6#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<7#0: <I64*32>>,), name='7')",
+            "Op(kind=OpKind.SetCA, "
+            "inputs=[], "
+            "immediates=[], "
+            "outputs=(<ca#0: <CA>>,), name='ca')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<8#0: <VL_MAXVL>>,), name='8')",
+            "Op(kind=OpKind.VecCopyToReg, "
+            "inputs=[<5#0: <I64*32>>, <8#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<9#0: <I64*32>>,), name='9')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<10#0: <VL_MAXVL>>,), name='10')",
+            "Op(kind=OpKind.VecCopyToReg, "
+            "inputs=[<7#0: <I64*32>>, <10#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<11#0: <I64*32>>,), name='11')",
+            "Op(kind=OpKind.SvAddE, "
+            "inputs=[<9#0: <I64*32>>, <11#0: <I64*32>>, <ca#0: <CA>>, "
+            "<vl#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<add#0: <I64*32>>, <add#1: <CA>>), name='add')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<12#0: <VL_MAXVL>>,), name='12')",
+            "Op(kind=OpKind.VecCopyFromReg, "
+            "inputs=[<add#0: <I64*32>>, <12#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<13#0: <I64*32>>,), name='13')",
+            "Op(kind=OpKind.SetVLI, "
+            "inputs=[], "
+            "immediates=[32], "
+            "outputs=(<14#0: <VL_MAXVL>>,), name='14')",
+            "Op(kind=OpKind.VecCopyToReg, "
+            "inputs=[<13#0: <I64*32>>, <14#0: <VL_MAXVL>>], "
+            "immediates=[], "
+            "outputs=(<15#0: <I64*32>>,), name='15')",
+            "Op(kind=OpKind.CopyToReg, "
+            "inputs=[<2#0: <I64>>], "
+            "immediates=[], "
+            "outputs=(<16#0: <I64>>,), name='16')",
+            "Op(kind=OpKind.SvStd, "
+            "inputs=[<15#0: <I64*32>>, <16#0: <I64>>, <vl#0: <VL_MAXVL>>], "
+            "immediates=[0], "
+            "outputs=(), name='st')",
+        ])
+
     def test_sim(self):
         fn, arg = self.make_add_fn()
         addr = 0x100
diff --git a/src/bigint_presentation_code/compiler_ir2.py b/src/bigint_presentation_code/compiler_ir2.py
index e6ffe4e..382d66c 100644
--- a/src/bigint_presentation_code/compiler_ir2.py
+++ b/src/bigint_presentation_code/compiler_ir2.py
@@ -56,6 +56,52 @@ class Fn:
         for op in self.ops:
             op.pre_ra_sim(state)
 
+    def pre_ra_insert_copies(self):
+        # type: () -> None
+        orig_ops = list(self.ops)
+        copied_outputs = {}  # type: dict[SSAVal, SSAVal]
+        self.ops.clear()
+        for op in orig_ops:
+            for i in range(len(op.inputs)):
+                inp = copied_outputs[op.inputs[i]]
+                if inp.ty.base_ty is BaseTy.I64:
+                    maxvl = inp.ty.reg_len
+                    if inp.ty.reg_len != 1:
+                        setvl = self.append_new_op(OpKind.SetVLI,
+                                                   immediates=[maxvl])
+                        vl = setvl.outputs[0]
+                        mv = self.append_new_op(OpKind.VecCopyToReg,
+                                                inputs=[inp, vl], maxvl=maxvl)
+                    else:
+                        mv = self.append_new_op(OpKind.CopyToReg, inputs=[inp])
+                    op.inputs[i] = mv.outputs[0]
+                elif inp.ty.base_ty is BaseTy.CA \
+                        or inp.ty.base_ty is BaseTy.VL_MAXVL:
+                    # all copies would be no-ops, so we don't need to copy
+                    op.inputs[i] = inp
+                else:
+                    assert_never(inp.ty.base_ty)
+            self.ops.append(op)
+            for out in op.outputs:
+                if out.ty.base_ty is BaseTy.I64:
+                    maxvl = out.ty.reg_len
+                    if out.ty.reg_len != 1:
+                        setvl = self.append_new_op(OpKind.SetVLI,
+                                                   immediates=[maxvl])
+                        vl = setvl.outputs[0]
+                        mv = self.append_new_op(OpKind.VecCopyFromReg,
+                                                inputs=[out, vl], maxvl=maxvl)
+                    else:
+                        mv = self.append_new_op(OpKind.CopyFromReg,
+                                                inputs=[out])
+                    copied_outputs[out] = mv.outputs[0]
+                elif out.ty.base_ty is BaseTy.CA \
+                        or out.ty.base_ty is BaseTy.VL_MAXVL:
+                    # all copies would be no-ops, so we don't need to copy
+                    copied_outputs[out] = out
+                else:
+                    assert_never(out.ty.base_ty)
+
 
 @unique
 @final
-- 
2.30.2