Add global, address_of, array_ref, load, store, ref<T>

Author: alastairreid

asl_xdsl/dialects/asl.py

@@ -5,6 +5,12 @@
 from typing import ClassVar, TypeAlias
 
 from xdsl.dialects import builtin
+from xdsl.dialects.builtin import (
+    IntAttr,
+    StringAttr,
+    SymbolRefAttr,
+    TupleType,
+)
 from xdsl.dialects.utils import parse_func_op_like, print_func_op_like
 from xdsl.ir import (
     Attribute,
@@ -18,8 +24,11 @@
     VerifyException,
 )
 from xdsl.irdl import (
+    AnyAttr,
     BaseAttr,
+    GenericAttrConstraint,
     IRDLOperation,
+    ParamAttrConstraint,
     ParameterDef,
     VarConstraint,
     irdl_attr_definition,
@@ -37,9 +46,11 @@
 from xdsl.printer import Printer
 from xdsl.traits import (
     CallableOpInterface,
+    ConstantLike,
     HasParent,
     IsolatedFromAbove,
     IsTerminator,
+    Pure,
     SymbolOpInterface,
     SymbolTable,
     SymbolUserOpInterface,
@@ -264,6 +275,24 @@ def print_parameters(self, printer: Printer) -> None:
             printer.print_string("x")
             printer.print_attribute(self.element_type)
 
+    @classmethod
+    def constr(
+        cls,
+        element_type: IRDLAttrConstraint | None = None,
+        *,
+        shape: IRDLGenericAttrConstraint[builtin.ArrayAttr[IntAttr]] | None = None,
+    ) -> GenericAttrConstraint[ArrayType]:
+        if element_type is None and shape is None:
+            return BaseAttr[ArrayType](ArrayType)
+        shape_constr = AnyAttr() if shape is None else shape
+        return ParamAttrConstraint[ArrayType](
+            ArrayType,
+            (
+                shape_constr,
+                element_type,
+            ),
+        )
+
 
 @irdl_attr_definition
 class BitVectorAttr(ParametrizedAttribute):
@@ -1558,6 +1587,163 @@ def __init__(
         )
 
 
+@irdl_attr_definition
+class ReferenceType(ParametrizedAttribute, TypeAttribute):
+    """
+    The type of a reference to an object
+    """
+
+    name = "asl.ref"
+    type: ParameterDef[Attribute]
+
+    def print_parameters(self, printer: Printer) -> None:
+        # We need this to pretty print a tuple and its members if
+        # this is referencing one, otherwise just let the type
+        # handle its own printing
+        printer.print("<")
+        printer.print(self.type)
+        printer.print(">")
+
+    @classmethod
+    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
+        # This is complicated by the fact we need to parse tuple
+        # here also as the buildin dialect does not support this
+        # yet
+        parser.parse_characters("<")
+        has_tuple = parser.parse_optional_keyword("tuple")
+        if has_tuple is None:
+            param_type = parser.parse_type()
+            parser.parse_characters(">")
+            return [param_type]
+        else:
+            # If its a tuple then there are any number of types
+            def parse_types():
+                return parser.parse_type()
+
+            param_types = parser.parse_comma_separated_list(
+                parser.Delimiter.ANGLE, parse_types
+            )
+            parser.parse_characters(">")
+            return [TupleType(param_types)]
+
+    @classmethod
+    def constr(
+        cls,
+        type: IRDLAttrConstraint | None = None,
+    ) -> GenericAttrConstraint[ReferenceType]:
+        if type is None:
+            return BaseAttr[ReferenceType](ReferenceType)
+        return ParamAttrConstraint[ReferenceType](
+            ReferenceType,
+            (type,),
+        )
+
+
+@irdl_op_definition
+class GlobalOp(IRDLOperation):
+    name = "asl.global"
+
+    assembly_format = "$sym_name `:` $global_type attr-dict"
+
+    global_type = prop_def(Attribute)
+    sym_name = prop_def(StringAttr)
+
+    traits = traits_def(SymbolOpInterface())
+
+    def __init__(
+        self,
+        global_type: Attribute,
+        sym_name: str | StringAttr,
+    ):
+        if isinstance(sym_name, str):
+            sym_name = StringAttr(sym_name)
+
+        props: dict[str, Attribute] = {
+            "global_type": global_type,
+            "sym_name": sym_name,
+        }
+
+        super().__init__(properties=props)
+
+
+@irdl_op_definition
+class AddressOfOp(IRDLOperation):
+    """
+    Convert a global reference to an SSA-value to be
+    used in other operations.
+
+    %p = asl.address_of @symbol : !asl.ref<i1>
+    """
+
+    name = "asl.address_of"
+
+    traits = traits_def(
+        ConstantLike(),
+        Pure(),
+    )
+
+    symbol = prop_def(SymbolRefAttr)
+    res = result_def()
+    
+    assembly_format = "`(` $symbol `)` `:` type($res) attr-dict"
+
+    assembly_format = "$symbol `:` type($res) attr-dict"
+
+
+@irdl_op_definition
+class ArrayRefOp(IRDLOperation):
+    """
+    Create a ref for an array element.
+
+    %element_ref = asl.array_ref %array_ref [ %index ]
+        : !asl.ref<!asl.array<16 x i8>> -> !asl.ref<i8>
+    """
+
+    name = "asl.array_ref"
+
+    traits = traits_def(
+        Pure(),
+    )
+
+    T: ClassVar = VarConstraint("T", AnyAttr())
+    A: ClassVar = VarConstraint("A", ArrayType.constr(T))
+    ref = operand_def(ReferenceType.constr(A))
+    index = operand_def(IntegerType())
+    res = result_def(ReferenceType.constr(T))
+
+    assembly_format = "$ref `[` $index `]` `:` type($ref) `->` type($res) attr-dict"
+
+
+@irdl_op_definition
+class LoadOp(IRDLOperation):
+    """
+    Load from a reference.
+    """
+
+    name = "asl.load"
+
+    T: ClassVar = VarConstraint("T", AnyAttr())
+    ref = operand_def(ReferenceType.constr(T))
+    res = result_def(T)
+
+    assembly_format = "`from` $ref `:` type($res) attr-dict"
+
+
+@irdl_op_definition
+class StoreOp(IRDLOperation):
+    """
+    Store from a reference.
+    """
+
+    name = "asl.store"
+
+    T: ClassVar = VarConstraint("T", AnyAttr())
+    ref = operand_def(ReferenceType.constr(T))
+    value = operand_def(T)
+
+    assembly_format = "$value `to` $ref `:` type($value) attr-dict"
+
+
 ASLDialect = Dialect(
     "asl",
     [
@@ -1622,12 +1808,19 @@ def __init__(
         # Slices
         GetSliceOp,
         SetSliceOp,
+        # References
+        GlobalOp,
+        AddressOfOp,
+        ArrayRefOp,
+        LoadOp,
+        StoreOp,
     ],
     [
         IntegerType,
         BitVectorType,
         BitVectorAttr,
         StringType,
         ArrayType,
+        ReferenceType,
     ],
 )

asl_xdsl/tools/asl_opt.py

@@ -34,14 +34,15 @@ def register_all_targets(self):
 
         def interpret_target(module: ModuleOp, output: IO[str]):
             from xdsl.interpreter import Interpreter
-            from xdsl.interpreters import arith, scf
+            from xdsl.interpreters import arith, cf, func
 
             from asl_xdsl.interpreters.asl import ASLFunctions
 
             interpreter = Interpreter(module, file=output)
             interpreter.register_implementations(ASLFunctions())
             interpreter.register_implementations(arith.ArithFunctions())
-            interpreter.register_implementations(scf.ScfFunctions())
+            interpreter.register_implementations(cf.CfFunctions())
+            interpreter.register_implementations(func.FuncFunctions())
             op = interpreter.get_op_for_symbol("main.0")
             trait = op.get_trait(CallableOpInterface)
             assert trait is not None

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "asl-xdsl"
 version = "0.0.0"
-dependencies = ["xdsl==0.40.0"]
+dependencies = ["xdsl==0.43.0"]
 requires-python = ">=3.10"
 
 [project.optional-dependencies]

tests/filecheck/dialects/asl/primitives.mlir

@@ -115,4 +115,21 @@ builtin.module {
     asl.print_sintN_dec %sint1 : i8
 // CHECK:         asl.print_sintN_hex %sint1 : i8 -> ()
 // CHECK-NEXT:    asl.print_sintN_dec %sint1 : i8 -> ()
+
+    asl.global "G" : !asl.bits<32>
+    asl.global "A" : !asl.array<16 x !asl.bits<32>>
+
+    %gref = asl.address_of @G : !asl.ref<!asl.bits<32>>
+    %aref = asl.address_of @A : !asl.ref<!asl.array<16 x !asl.bits<32>>>
+    %eref = asl.array_ref %aref[%int1] : !asl.ref<!asl.array<16 x !asl.bits<32>>> -> !asl.ref<!asl.bits<32>>
+    asl.store %bits1 to %gref : !asl.bits<32>
+    %load = asl.load from %eref : !asl.bits<32>
+// CHECK:         asl.global "G" : !asl.bits<32>
+// CHECK-NEXT:    asl.global "A" : !asl.array<16x!asl.bits<32>>
+// CHECK-NEXT:    %gref = asl.address_of @G : !asl.ref<!asl.bits<32>>
+// CHECK-NEXT:    %aref = asl.address_of @A : !asl.ref<!asl.array<16x!asl.bits<32>>>
+// CHECK-NEXT:    %eref = asl.array_ref %aref[%int1] : !asl.ref<!asl.array<16x!asl.bits<32>>> -> !asl.ref<!asl.bits<32>>
+// CHECK-NEXT:    asl.store %bits1 to %gref : !asl.bits<32>
+// CHECK-NEXT:    %load = asl.load from %eref : !asl.bits<32>
+
 }