From 4a026317a8b26b3e38bc8cb6033f1ef59e8a961d Mon Sep 17 00:00:00 2001
From: emmau678 <emmau678@gmail.com>
Date: Mon, 25 Nov 2024 14:02:12 +0000
Subject: [PATCH] dialects (arm): add mul instruction with 2 src regs

---
 tests/filecheck/dialects/arm/test_ops.mlir |  8 +++++
 xdsl/dialects/arm/__init__.py              |  3 +-
 xdsl/dialects/arm/ops.py                   | 40 ++++++++++++++++++++++
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/tests/filecheck/dialects/arm/test_ops.mlir b/tests/filecheck/dialects/arm/test_ops.mlir
index 4ce5afbf2b..3a2020293b 100644
--- a/tests/filecheck/dialects/arm/test_ops.mlir
+++ b/tests/filecheck/dialects/arm/test_ops.mlir
@@ -6,9 +6,17 @@
 // 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: %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: %dss_mul = arm.dss.mul %x1, %x2 {"comment" = "multiply s1 by s2"} : (!arm.reg<x1>, !arm.reg<x2>) -> !arm.reg<x3>
+// CHECK-ASM: mul x3, x1, x2 # multiply s1 by s2
+%dss_mul = arm.dss.mul %x1, %x2 {"comment" = "multiply s1 by s2"} : (!arm.reg<x1>, !arm.reg<x2>) -> !arm.reg<x3>
+
 // 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>
+// CHECK-GENERIC: %dss_mul = "arm.dss.mul"(%x1, %x2) {"comment" = "multiply s1 by s2"} : (!arm.reg<x1>, !arm.reg<x2>) -> !arm.reg<x3>
diff --git a/xdsl/dialects/arm/__init__.py b/xdsl/dialects/arm/__init__.py
index 0c6bfa6cba..dbb2afbdba 100644
--- a/xdsl/dialects/arm/__init__.py
+++ b/xdsl/dialects/arm/__init__.py
@@ -8,7 +8,7 @@
 from xdsl.dialects.builtin import ModuleOp
 from xdsl.ir import Dialect
 
-from .ops import ARMOperation, DSMovOp, GetRegisterOp
+from .ops import ARMOperation, DSMovOp, DSSMulOp, GetRegisterOp
 from .register import IntRegisterType
 
 
@@ -25,6 +25,7 @@ def print_assembly(module: ModuleOp, output: IO[str]) -> None:
     [
         GetRegisterOp,
         DSMovOp,
+        DSSMulOp,
     ],
     [
         IntRegisterType,
diff --git a/xdsl/dialects/arm/ops.py b/xdsl/dialects/arm/ops.py
index 97ea772c1f..9c152d82b2 100644
--- a/xdsl/dialects/arm/ops.py
+++ b/xdsl/dialects/arm/ops.py
@@ -114,3 +114,43 @@ def __init__(self, register_type: IntRegisterType):
 
     def assembly_line(self):
         return None
+
+
+@irdl_op_definition
+class DSSMulOp(ARMInstruction):
+    """
+    Multiplies the values in s1 and s2 and stores the result in d.
+
+    https://developer.arm.com/documentation/ddi0597/2024-06/Base-Instructions/MUL--MULS--Multiply-?lang=en
+    """
+
+    name = "arm.dss.mul"
+
+    d = result_def(IntRegisterType)
+    s1 = operand_def(IntRegisterType)
+    s2 = operand_def(IntRegisterType)
+    assembly_format = (
+        "$s1 `,` $s2 attr-dict `:` `(` type($s1) `,` type($s2) `)` `->` type($d)"
+    )
+
+    def __init__(
+        self,
+        d: IntRegisterType,
+        s1: Operation | SSAValue,
+        s2: Operation | SSAValue,
+        *,
+        comment: str | StringAttr | None = None,
+    ):
+        if isinstance(comment, str):
+            comment = StringAttr(comment)
+
+        super().__init__(
+            operands=(s1, s2),
+            attributes={
+                "comment": comment,
+            },
+            result_types=(d,),
+        )
+
+    def assembly_line_args(self):
+        return (self.d, self.s1, self.s2)