From 95ba97bd120ab39a62d695e64bbc2b3d3e576db7 Mon Sep 17 00:00:00 2001
From: emmau678 <emmau678@gmail.com>
Date: Fri, 17 Jan 2025 13:40:04 +0000
Subject: [PATCH] [WIP]: Initialise arm_cf and begin beq implementation

---
 tests/filecheck/dialects/arm_cf/test_ops.mlir | 19 ++++++
 tests/test_assembly_arg_str.py                |  5 ++
 xdsl/dialects/__init__.py                     |  6 ++
 xdsl/dialects/arm/assembly.py                 |  4 +-
 xdsl/dialects/arm_cf.py                       | 65 +++++++++++++++++++
 5 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 tests/filecheck/dialects/arm_cf/test_ops.mlir
 create mode 100644 xdsl/dialects/arm_cf.py

diff --git a/tests/filecheck/dialects/arm_cf/test_ops.mlir b/tests/filecheck/dialects/arm_cf/test_ops.mlir
new file mode 100644
index 0000000000..6b038d6bed
--- /dev/null
+++ b/tests/filecheck/dialects/arm_cf/test_ops.mlir
@@ -0,0 +1,19 @@
+// RUN: XDSL_ROUNDTRIP
+// RUN: XDSL_GENERIC_ROUNDTRIP
+// RUN: xdsl-opt -t arm-asm %s | filecheck %s --check-prefix=CHECK-ASM
+
+// CHECK: %x1 = arm.get_register : !arm.reg<x1>
+%x1 = arm.get_register : !arm.reg<x1>
+
+// CHECK: %x2 = arm.get_register : !arm.reg<x2>
+%x2 = arm.get_register : !arm.reg<x2>
+
+// CHECK: arm.label "testlabel" {comment = "this is a label"}
+// CHECK-ASM: testlabel:                                       # this is a label
+arm.label "testlabel" {comment = "this is a label"}
+
+// CHECK-NEXT: arm_cf.beq %x1, %x2, ^testlabel
+// CHECK-ASM: arm_cf.beq x1, x2, ^testlabel
+arm_cf.beq %x1, %x2, ^testlabel : (!arm.reg<x1>, !arm.reg<x2>)
+
+// CHECK-GENERIC: "arm.label"() <{label = "testlabel"}> {comment = "this is a label"} : () -> ()
diff --git a/tests/test_assembly_arg_str.py b/tests/test_assembly_arg_str.py
index 8552b2c805..9896616d56 100644
--- a/tests/test_assembly_arg_str.py
+++ b/tests/test_assembly_arg_str.py
@@ -11,6 +11,11 @@ def test_assembly_arg_str_ARMRegister():
     assert assembly_arg_str(arg) == "x0"
 
 
+def test_assembly_arg_str_str():
+    arg = "string_arg"
+    assert assembly_arg_str(arg) == "string_arg"
+
+
 def test_assembly_arg_str_SSAValue_valid():
     arg = TestSSAValue(IntRegisterType("x1"))
     assert assembly_arg_str(arg) == "x1"
diff --git a/xdsl/dialects/__init__.py b/xdsl/dialects/__init__.py
index 0c8eb9e589..c397b1db12 100644
--- a/xdsl/dialects/__init__.py
+++ b/xdsl/dialects/__init__.py
@@ -33,6 +33,11 @@ def get_arm():
 
         return ARM
 
+    def get_arm_cf():
+        from xdsl.dialects.arm_cf import ARM_CF
+
+        return ARM_CF
+
     def get_arm_func():
         from xdsl.dialects.arm_func import ARM_FUNC
 
@@ -329,6 +334,7 @@ def get_transform():
         "air": get_air,
         "arith": get_arith,
         "arm": get_arm,
+        "arm_cf": get_arm_cf,
         "arm_func": get_arm_func,
         "bufferization": get_bufferization,
         "builtin": get_builtin,
diff --git a/xdsl/dialects/arm/assembly.py b/xdsl/dialects/arm/assembly.py
index 2ae9e20380..b3fcec8942 100644
--- a/xdsl/dialects/arm/assembly.py
+++ b/xdsl/dialects/arm/assembly.py
@@ -4,7 +4,7 @@
 from xdsl.dialects.builtin import StringAttr
 from xdsl.ir import SSAValue
 
-AssemblyInstructionArg: TypeAlias = ARMRegisterType | SSAValue
+AssemblyInstructionArg: TypeAlias = ARMRegisterType | SSAValue | str
 
 
 def append_comment(line: str, comment: StringAttr | None) -> str:
@@ -20,6 +20,8 @@ def assembly_arg_str(arg: AssemblyInstructionArg) -> str:
     if isinstance(arg, ARMRegisterType):
         reg = arg.register_name
         return reg
+    elif isinstance(arg, str):
+        return arg
     else:  # SSAValue
         if isinstance(arg.type, ARMRegisterType):
             reg = arg.type.register_name
diff --git a/xdsl/dialects/arm_cf.py b/xdsl/dialects/arm_cf.py
new file mode 100644
index 0000000000..78133a8504
--- /dev/null
+++ b/xdsl/dialects/arm_cf.py
@@ -0,0 +1,65 @@
+from xdsl.dialects import arm
+from xdsl.dialects.builtin import StringAttr
+from xdsl.interpreter import Successor
+from xdsl.ir import Dialect, Operation, SSAValue
+from xdsl.irdl import (
+    irdl_op_definition,
+    operand_def,
+    successor_def,
+    traits_def,
+)
+from xdsl.traits import IsTerminator
+
+
+@irdl_op_definition
+class BEqOp(arm.ops.ARMInstruction):
+    """
+    Branch if equal
+    https://developer.arm.com/documentation/den0042/a/Unified-Assembly-Language-Instructions/Instruction-set-basics/Conditional-execution
+    """
+
+    name = "arm_cf.beq"
+    s1 = operand_def(arm.register.IntRegisterType)
+    s2 = operand_def(arm.register.IntRegisterType)
+
+    then_block = successor_def()
+
+    traits = traits_def(
+        IsTerminator(),
+    )
+
+    assembly_format = (
+        "$s1 `,` $s2 `,` $then_block attr-dict `:` `(` type($s1) `,` type($s2) `)`"
+    )
+
+    def __init__(
+        self,
+        s1: Operation | SSAValue,
+        s2: Operation | SSAValue,
+        then_block: Successor,
+        *,
+        comment: str | StringAttr | None = None,
+    ):
+        if isinstance(comment, str):
+            comment = StringAttr(comment)
+
+        super().__init__(
+            operands=[s1, s2],
+            attributes={
+                "comment": comment,
+            },
+            successors=(then_block),
+        )
+
+    def assembly_line_args(self):
+        then_label = self.then_block.first_op
+        assert isinstance(then_label, arm.ops.LabelOp)
+        return (self.s1, self.s2, then_label.label.data)
+
+
+ARM_CF = Dialect(
+    "arm_cf",
+    [
+        BEqOp,
+    ],
+)