Repository: xdslproject/xdsl

Issue:
CSR instructions don't use new ABI constants to verify x0 register

Discussion:

Solution patch:
diff --git a/xdsl/dialects/arith.py b/xdsl/dialects/arith.py
index 75daa4dff..cea60c46d 100644
--- a/xdsl/dialects/arith.py
+++ b/xdsl/dialects/arith.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import abc
 from collections.abc import Mapping, Sequence
-from typing import ClassVar, Literal, TypeVar, cast, overload
+from typing import ClassVar, Literal, TypeVar, cast
 
 from xdsl.dialects.builtin import (
     AnyFloat,
@@ -146,24 +146,14 @@ class ConstantOp(IRDLOperation):
 
     assembly_format = "attr-dict $value"
 
-    @overload
     def __init__(
         self,
         value: IntegerAttr | FloatAttr[AnyFloat] | DenseIntOrFPElementsAttr,
-        value_type: None = None,
-    ) -> None: ...
-
-    @overload
-    def __init__(self, value: Attribute, value_type: Attribute) -> None: ...
-
-    def __init__(
-        self,
-        value: IntegerAttr | FloatAttr[AnyFloat] | Attribute,
         value_type: Attribute | None = None,
     ):
         if value_type is None:
-            value = cast(IntegerAttr | FloatAttr[AnyFloat], value)
-            value_type = value.type
+            value_type = value.get_type()
+
         super().__init__(
             operands=[], result_types=[value_type], properties={"value": value}
         )
diff --git a/xdsl/dialects/riscv.py b/xdsl/dialects/riscv.py
index 72c5223ee..8b621f3f3 100644
--- a/xdsl/dialects/riscv.py
+++ b/xdsl/dialects/riscv.py
@@ -19,6 +19,7 @@ from xdsl.dialects.builtin import (
     IntegerAttr,
     IntegerType,
     ModuleOp,
+    NoneAttr,
     Signedness,
     StringAttr,
     UnitAttr,
@@ -67,6 +68,15 @@ from xdsl.utils.exceptions import VerifyException
 from xdsl.utils.hints import isa
 
 
+def is_non_zero(reg: IntRegisterType) -> bool:
+    """
+    Returns True if the register is allocated, and is not the x0/ZERO register.
+    """
+    return (
+        reg.is_allocated and not isinstance(reg.index, NoneAttr) and reg.index.data != 0
+    )
+
+
 @irdl_attr_definition
 class FastMathFlagsAttr(FastMathAttrBase):
     """
@@ -1181,9 +1191,8 @@ class CsrReadWriteOperation(RISCVCustomFormatOperation, RISCVInstruction, ABC):
     def verify_(self) -> None:
         if not self.writeonly:
             return
-        if not isinstance(self.rd.type, IntRegisterType):
-            return
-        if self.rd.type.is_allocated and self.rd.type != Registers.ZERO:
+        assert isinstance(self.rd.type, IntRegisterType)
+        if is_non_zero(self.rd.type):
             raise VerifyException(
                 "When in 'writeonly' mode, destination must be register x0 (a.k.a. 'zero'), "
                 f"not '{self.rd.type.register_name.data}'"
@@ -1255,9 +1264,8 @@ class CsrBitwiseOperation(RISCVCustomFormatOperation, RISCVInstruction, ABC):
     def verify_(self) -> None:
         if not self.readonly:
             return
-        if not isinstance(self.rs1.type, IntRegisterType):
-            return
-        if self.rs1.type.is_allocated and self.rs1.type != Registers.ZERO:
+        assert isinstance(self.rs1.type, IntRegisterType)
+        if is_non_zero(self.rs1.type):
             raise VerifyException(
                 "When in 'readonly' mode, source must be register x0 (a.k.a. 'zero'), "
                 f"not '{self.rs1.type.register_name.data}'"
@@ -1327,9 +1335,8 @@ class CsrReadWriteImmOperation(RISCVCustomFormatOperation, RISCVInstruction, ABC
     def verify_(self) -> None:
         if self.writeonly is None:
             return
-        if not isinstance(self.rd.type, IntRegisterType):
-            return
-        if self.rd.type.is_allocated and self.rd.type != Registers.ZERO:
+        assert isinstance(self.rd.type, IntRegisterType)
+        if is_non_zero(self.rd.type):
             raise VerifyException(
                 "When in 'writeonly' mode, destination must be register x0 (a.k.a. 'zero'), "
                 f"not '{self.rd.type.register_name.data}'"
diff --git a/xdsl/transforms/experimental/function_constant_pinning.py b/xdsl/transforms/experimental/function_constant_pinning.py
index a405a20c2..d9ccf102a 100644
--- a/xdsl/transforms/experimental/function_constant_pinning.py
+++ b/xdsl/transforms/experimental/function_constant_pinning.py
@@ -3,8 +3,14 @@ from typing import cast
 
 from xdsl.context import Context
 from xdsl.dialects import arith, builtin, func, scf
-from xdsl.dialects.builtin import ArrayAttr, StringAttr
-from xdsl.ir import Attribute, Block, Operation, Region
+from xdsl.dialects.builtin import (
+    ArrayAttr,
+    IndexType,
+    IntegerAttr,
+    IntegerType,
+    StringAttr,
+)
+from xdsl.ir import Block, Operation, Region
 from xdsl.passes import ModulePass
 from xdsl.pattern_rewriter import (
     PatternRewriter,
@@ -113,7 +119,7 @@ class FunctionConstantPinning(RewritePattern):
 
 def generate_func_with_pinned_val(
     func_op: func.FuncOp,
-    pin: Attribute,
+    pin: IntegerAttr[IntegerType | IndexType],
     rewriter: PatternRewriter,
 ):
     """
@@ -171,7 +177,9 @@ def func_contains_pinning_annotation(funcop: func.FuncOp) -> Operation | None:
             return op
 
 
-def get_pinned_vals_for_op(op: Operation) -> list[Attribute] | None:
+def get_pinned_vals_for_op(
+    op: Operation,
+) -> list[IntegerAttr[IntegerType | IndexType]] | None:
     """
     Reads the "pin_to_constants" attribute of an operation, checks for valid
     formatting, and return the list of attribute values that should be pinned.
@@ -182,7 +190,7 @@ def get_pinned_vals_for_op(op: Operation) -> list[Attribute] | None:
     if not isinstance(pin_attr, ArrayAttr):
         return None
 
-    return list(cast(ArrayAttr[Attribute], pin_attr))
+    return list(cast(ArrayAttr[IntegerAttr[IntegerType | IndexType]], pin_attr))
 
 
 def ops_between_op_and_func_start(


Test patch:
diff --git a/tests/dialects/test_arith.py b/tests/dialects/test_arith.py
index 0e72fefef..df64d0028 100644
--- a/tests/dialects/test_arith.py
+++ b/tests/dialects/test_arith.py
@@ -52,6 +52,7 @@ from xdsl.dialects.arith import (
 from xdsl.dialects.builtin import (
     AnyTensorType,
     AnyVectorType,
+    DenseIntOrFPElementsAttr,
     FloatAttr,
     IndexType,
     IntegerAttr,
@@ -128,6 +129,18 @@ class Test_integer_arith_construction:
         _ = CmpiOp(self.a, self.b, input)
 
 
+def test_constant_construction():
+    c1 = ConstantOp(IntegerAttr(1, i32))
+    assert c1.value.type == i32
+
+    c3 = ConstantOp(FloatAttr(1.0, f32))
+    assert c3.value.type == f32
+
+    value_type = TensorType(i32, [2, 2])
+    c5 = ConstantOp(DenseIntOrFPElementsAttr.create_dense_int(value_type, [1, 2, 3, 4]))
+    assert c5.value.type == value_type
+
+
 @pytest.mark.parametrize(
     "value, truncated",
     [
diff --git a/tests/dialects/test_riscv.py b/tests/dialects/test_riscv.py
index 64d68a41f..59a055e5d 100644
--- a/tests/dialects/test_riscv.py
+++ b/tests/dialects/test_riscv.py
@@ -38,6 +38,23 @@ def test_add_op():
     assert isinstance(riscv.IntRegisterType.infinite_register(1).index, NoneAttr)
 
 
+def test_is_non_zero():
+    # Test zero register
+    x0_reg = riscv.IntRegisterType.from_name("x0")
+    assert not riscv.is_non_zero(riscv.Registers.ZERO)
+    assert not riscv.is_non_zero(x0_reg)
+
+    # Test non-zero registers
+    a0_reg = riscv.Registers.A0
+    t0_reg = riscv.Registers.T0
+    assert riscv.is_non_zero(a0_reg)
+    assert riscv.is_non_zero(t0_reg)
+
+    # Test unallocated register
+    unalloc_reg = riscv.IntRegisterType.unallocated()
+    assert not riscv.is_non_zero(unalloc_reg)
+
+
 def test_csr_op():
     a1 = TestSSAValue(riscv.Registers.A1)
     zero = TestSSAValue(riscv.Registers.ZERO)