dialects (arm): add mov op (#3476)

Add a mov op that moves data from source to destination

---------

Co-authored-by: emmau678 <eu233@Emma-laptop>
Co-authored-by: Sasha Lopoukhine <[email protected]>

tests/filecheck/dialects/arm/test_ops.mlir

@@ -1,9 +1,14 @@
 // 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: %ds_mov = arm.ds.mov %x1 {"comment" = "move contents of s to d"} : (!arm.reg<x1>) -> !arm.reg<x2>
+// CHECK-ASM: mov x2, x1 # move contents of s to d
+%ds_mov = arm.ds.mov %x1 {"comment" = "move contents of s to d"} : (!arm.reg<x1>) -> !arm.reg<x2>
 
 // CHECK-GENERIC: %x1 = "arm.get_register"() : () -> !arm.reg<x1>
+// CHECK-GENERIC: %ds_mov = "arm.ds.mov"(%x1) {"comment" = "move contents of s to d"} : (!arm.reg<x1>) -> !arm.reg<x2>

xdsl/dialects/arm/__init__.py

@@ -3,15 +3,28 @@
 https://developer.arm.com/documentation/102374/0101/Overview
 """
 
+from typing import IO
+
+from xdsl.dialects.builtin import ModuleOp
 from xdsl.ir import Dialect
 
-from .ops import GetRegisterOp
+from .ops import ARMOperation, DSMovOp, GetRegisterOp
 from .register import IntRegisterType
 
+
+def print_assembly(module: ModuleOp, output: IO[str]) -> None:
+    for op in module.body.walk():
+        assert isinstance(op, ARMOperation), f"{op}"
+        asm = op.assembly_line()
+        if asm is not None:
+            print(asm, file=output)
+
+
 ARM = Dialect(
     "arm",
     [
         GetRegisterOp,
+        DSMovOp,
     ],
     [
         IntRegisterType,

xdsl/dialects/arm/assembly.py

@@ -0,0 +1,38 @@
+from typing import TypeAlias
+
+from xdsl.dialects.arm.register import ARMRegisterType
+from xdsl.dialects.builtin import StringAttr
+from xdsl.ir import SSAValue
+
+AssemblyInstructionArg: TypeAlias = SSAValue
+
+
+def append_comment(line: str, comment: StringAttr | None) -> str:
+    if comment is None:
+        return line
+
+    padding = " " * max(0, 48 - len(line))
+
+    return f"{line}{padding} # {comment.data}"
+
+
+def assembly_arg_str(arg: AssemblyInstructionArg) -> str:
+    if isinstance(arg.type, ARMRegisterType):
+        reg = arg.type.register_name
+        return reg
+    else:
+        raise ValueError(f"Unexpected register type {arg.type}")
+
+
+def assembly_line(
+    name: str,
+    arg_str: str,
+    comment: StringAttr | None = None,
+    is_indented: bool = True,
+) -> str:
+    code = "    " if is_indented else ""
+    code += name
+    if arg_str:
+        code += f" {arg_str}"
+    code = append_comment(code, comment)
+    return code

xdsl/dialects/arm/ops.py

@@ -1,7 +1,16 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 
-from xdsl.irdl import IRDLOperation, irdl_op_definition, result_def
+from xdsl.dialects.builtin import StringAttr
+from xdsl.ir import Operation, SSAValue
+from xdsl.irdl import (
+    IRDLOperation,
+    irdl_op_definition,
+    operand_def,
+    opt_attr_def,
+    result_def,
+)
 
+from .assembly import AssemblyInstructionArg, assembly_arg_str, assembly_line
 from .register import IntRegisterType
 
 
@@ -10,7 +19,83 @@ class ARMOperation(IRDLOperation, ABC):
     Base class for operations that can be a part of ARM assembly printing.
     """
 
-    ...
+    @abstractmethod
+    def assembly_line(self) -> str | None:
+        raise NotImplementedError()
+
+
+class ARMInstruction(ARMOperation, ABC):
+    """
+    Base class for operations that can be a part of x86 assembly printing. Must
+    represent an instruction in the x86 instruction set.
+    The name of the operation will be used as the x86 assembly instruction name.
+    """
+
+    comment = opt_attr_def(StringAttr)
+    """
+    An optional comment that will be printed along with the instruction.
+    """
+
+    @abstractmethod
+    def assembly_line_args(self) -> tuple[AssemblyInstructionArg | None, ...]:
+        """
+        The arguments to the instruction, in the order they should be printed in the
+        assembly.
+        """
+        raise NotImplementedError()
+
+    def assembly_instruction_name(self) -> str:
+        """
+        By default, the name of the instruction is the same as the name of the operation.
+        """
+
+        return self.name.split(".")[-1]
+
+    def assembly_line(self) -> str | None:
+        # default assembly code generator
+        instruction_name = self.assembly_instruction_name()
+        arg_str = ", ".join(
+            assembly_arg_str(arg)
+            for arg in self.assembly_line_args()
+            if arg is not None
+        )
+        return assembly_line(instruction_name, arg_str, self.comment)
+
+
+@irdl_op_definition
+class DSMovOp(ARMInstruction):
+    """
+    Copies the value of s into d.
+
+    https://developer.arm.com/documentation/dui0473/m/arm-and-thumb-instructions/mov
+    """
+
+    name = "arm.ds.mov"
+
+    d = result_def(IntRegisterType)
+    s = operand_def(IntRegisterType)
+    assembly_format = "$s attr-dict `:` `(` type($s) `)` `->` type($d)"
+
+    def __init__(
+        self,
+        d: IntRegisterType,
+        s: Operation | SSAValue,
+        *,
+        comment: str | StringAttr | None = None,
+    ):
+        if isinstance(comment, str):
+            comment = StringAttr(comment)
+
+        super().__init__(
+            operands=(s,),
+            attributes={
+                "comment": comment,
+            },
+            result_types=(d,),
+        )
+
+    def assembly_line_args(self):
+        return (self.d, self.s)
 
 
 @irdl_op_definition
@@ -26,3 +111,6 @@ class GetRegisterOp(ARMOperation):
 
     def __init__(self, register_type: IntRegisterType):
         super().__init__(result_types=[register_type])
+
+    def assembly_line(self):
+        return None

xdsl/xdsl_opt_main.py

@@ -192,6 +192,11 @@ def register_all_targets(self):
         Add other/additional targets by overloading this function.
         """
 
+        def _output_arm_asm(prog: ModuleOp, output: IO[str]):
+            from xdsl.dialects.arm import print_assembly
+
+            print_assembly(prog, output)
+
         def _output_mlir(prog: ModuleOp, output: IO[str]):
             printer = Printer(
                 stream=output,
@@ -241,6 +246,7 @@ def _print_to_csl(prog: ModuleOp, output: IO[str]):
 
             print_to_csl(prog, output)
 
+        self.available_targets["arm-asm"] = _output_arm_asm
         self.available_targets["mlir"] = _output_mlir
         self.available_targets["riscv-asm"] = _output_riscv_asm
         self.available_targets["x86-asm"] = _output_x86_asm