core: Implement EffectKind

tests/filecheck/transforms/cse.mlir

@@ -313,15 +313,13 @@ func.func @use_before_def() {
 
 // CHECK:         func.func @remove_direct_duplicated_read_op() -> i32 {
 // CHECK-NEXT:      %0 = "test.op_with_memread"() : () -> i32
-// CHECK-NEXT:      %1 = "test.op_with_memread"() : () -> i32
-// CHECK-NEXT:      %2 = arith.addi %0, %1 : i32
-// CHECK-NEXT:      func.return %2 : i32
+// CHECK-NEXT:      %1 = arith.addi %0, %0 : i32
+// CHECK-NEXT:      func.return %1 : i32
 // CHECK-NEXT:    }
 
 
 /// This test is checking that CSE is removing duplicated read op that follow
 /// other.
-/// NB: xDSL doesn't, we don't have the notion of "read" ops.
 // CHECK-LABEL: @remove_multiple_duplicated_read_op
   func.func @remove_multiple_duplicated_read_op() -> i64 {
     %55 = "test.op_with_memread"() : () -> i64
@@ -336,19 +334,14 @@ func.func @use_before_def() {
 
 // CHECK:         func.func @remove_multiple_duplicated_read_op() -> i64 {
 // CHECK-NEXT:      %0 = "test.op_with_memread"() : () -> i64
-// CHECK-NEXT:      %1 = "test.op_with_memread"() : () -> i64
-// CHECK-NEXT:      %2 = arith.addi %0, %1 : i64
-// CHECK-NEXT:      %3 = "test.op_with_memread"() : () -> i64
-// CHECK-NEXT:      %4 = arith.addi %2, %3 : i64
-// CHECK-NEXT:      %5 = "test.op_with_memread"() : () -> i64
-// CHECK-NEXT:      %6 = arith.addi %4, %5 : i64
-// CHECK-NEXT:      func.return %6 : i64
+// CHECK-NEXT:      %1 = arith.addi %0, %0 : i64
+// CHECK-NEXT:      %2 = arith.addi %1, %0 : i64
+// CHECK-NEXT:      %3 = arith.addi %2, %0 : i64
+// CHECK-NEXT:      func.return %3 : i64
 // CHECK-NEXT:    }
 
 /// This test is checking that CSE is not removing duplicated read op that
 /// have write op in between.
-/// NB: xDSL doesn't, we don't have the notion of "read" ops.
-// CHECK-LABEL: @dont_remove_duplicated_read_op_with_sideeffecting
 func.func @dont_remove_duplicated_read_op_with_sideeffecting() -> i32 {
     %62 = "test.op_with_memread"() : () -> i32
     "test.op_with_memwrite"() : () -> ()
@@ -588,7 +581,6 @@ func.func @no_cse_multiple_regions_side_effect(%arg0_12 : i1, %arg1_9 : memref<5
 // CHECK-NEXT:      func.return %0, %2 : memref<5xf32>, memref<5xf32>
 // CHECK-NEXT:    }
 
-// xDSL doesn't have the notion of sideffects.
  func.func @cse_recursive_effects_success() -> (i32, i32, i32) {
     %98 = "test.op_with_memread"() : () -> i32
     %99 = arith.constant true
@@ -613,8 +605,7 @@ func.func @no_cse_multiple_regions_side_effect(%arg0_12 : i1, %arg1_9 : memref<5
 // CHECK-NEXT:        %4 = arith.constant 24 : i32
 // CHECK-NEXT:        scf.yield %4 : i32
 // CHECK-NEXT:      }) : (i1) -> i32
-// CHECK-NEXT:      %5 = "test.op_with_memread"() : () -> i32
-// CHECK-NEXT:      func.return %0, %5, %2 : i32, i32, i32
+// CHECK-NEXT:      func.return %0, %0, %2 : i32, i32, i32
 // CHECK-NEXT:    }
 
 // xDSL doesn't have the notion of sideffects.

xdsl/dialects/riscv.py

@@ -56,6 +56,7 @@
     IsolatedFromAbove,
     IsTerminator,
     MemoryEffect,
+    MemoryEffectKind,
     NoTerminator,
     Pure,
 )
@@ -2809,11 +2810,19 @@ class RegisterAllocatedMemoryEffect(MemoryEffect):
     """
 
     @classmethod
-    def has_effects(cls, op: Operation) -> bool:
-        return any(
+    def get_effects(cls, op: Operation) -> set[MemoryEffectKind]:
+        effects = set[MemoryEffectKind]()
+        if any(
             isinstance(r.type, RegisterType) and r.type.is_allocated
-            for r in chain(op.results, op.operands)
-        )
+            for r in chain(op.results)
+        ):
+            effects.add(MemoryEffectKind.WRITE)
+        if any(
+            isinstance(r.type, RegisterType) and r.type.is_allocated
+            for r in chain(op.operands)
+        ):
+            effects.add(MemoryEffectKind.READ)
+        return effects
 
 
 class GetAnyRegisterOperation(Generic[RDInvT], IRDLOperation, RISCVOp):

xdsl/dialects/test.py

@@ -28,7 +28,7 @@
 )
 from xdsl.parser import AttrParser
 from xdsl.printer import Printer
-from xdsl.traits import IsTerminator, Pure
+from xdsl.traits import IsTerminator, MemoryReadEffect, MemoryWriteEffect, Pure
 
 
 @irdl_op_definition
@@ -153,6 +153,92 @@ def __init__(
         )
 
 
+@irdl_op_definition
+class TestReadOp(IRDLOperation):
+    """
+    This operation can produce an arbitrary number of SSAValues with arbitrary
+    types. It is used in filecheck testing to reduce to artificial dependencies
+    on other dialects (i.e. dependencies that only come from the structure of
+    the test rather than the actual dialect).
+    Its main difference from TestOp is that it satisfies the MemoryReadEffect trait,
+    so we can test CSE - this op assumes it always and only reads.
+    """
+
+    name = "test.op_with_memread"
+
+    res: VarOpResult = var_result_def()
+    ops: VarOperand = var_operand_def()
+    regs: VarRegion = var_region_def()
+    successor: VarSuccessor = var_successor_def()
+
+    prop1 = opt_prop_def(Attribute)
+    prop2 = opt_prop_def(Attribute)
+    prop3 = opt_prop_def(Attribute)
+
+    traits = frozenset([MemoryReadEffect()])
+
+    def __init__(
+        self,
+        operands: Sequence[SSAValue | Operation] = (),
+        result_types: Sequence[Attribute] = (),
+        attributes: Mapping[str, Attribute | None] | None = None,
+        properties: Mapping[str, Attribute | None] | None = None,
+        successors: Sequence[Block] = (),
+        regions: Sequence[Region | Sequence[Operation] | Sequence[Block]] = (),
+    ):
+        super().__init__(
+            operands=(operands,),
+            result_types=(result_types,),
+            attributes=attributes,
+            properties=properties,
+            successors=(successors,),
+            regions=(regions,),
+        )
+
+
+@irdl_op_definition
+class TestWriteOp(IRDLOperation):
+    """
+    This operation can produce an arbitrary number of SSAValues with arbitrary
+    types. It is used in filecheck testing to reduce to artificial dependencies
+    on other dialects (i.e. dependencies that only come from the structure of
+    the test rather than the actual dialect).
+    Its main difference from TestOp is that it satisfies the MemoryReadEffect trait,
+    so we can test CSE - this op assumes it always and only writes.
+    """
+
+    name = "test.op_with_memwrite"
+
+    res: VarOpResult = var_result_def()
+    ops: VarOperand = var_operand_def()
+    regs: VarRegion = var_region_def()
+    successor: VarSuccessor = var_successor_def()
+
+    prop1 = opt_prop_def(Attribute)
+    prop2 = opt_prop_def(Attribute)
+    prop3 = opt_prop_def(Attribute)
+
+    traits = frozenset([MemoryWriteEffect()])
+
+    def __init__(
+        self,
+        operands: Sequence[SSAValue | Operation] = (),
+        result_types: Sequence[Attribute] = (),
+        attributes: Mapping[str, Attribute | None] | None = None,
+        properties: Mapping[str, Attribute | None] | None = None,
+        successors: Sequence[Block] = (),
+        regions: Sequence[Region | Sequence[Operation] | Sequence[Block]] = (),
+    ):
+        super().__init__(
+            operands=(operands,),
+            result_types=(result_types,),
+            attributes=attributes,
+            properties=properties,
+            successors=(successors,),
+            regions=(regions,),
+        )
+
+
 @irdl_attr_definition
 class TestType(Data[str], TypeAttribute):
     """
@@ -178,7 +264,9 @@ def print_parameter(self, printer: Printer) -> None:
     [
         TestOp,
         TestPureOp,
+        TestReadOp,
         TestTermOp,
+        TestWriteOp,
     ],
     [
         TestType,

xdsl/traits.py

@@ -3,7 +3,8 @@
 import abc
 from collections.abc import Iterator
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, TypeVar
+from enum import Enum, auto
+from typing import TYPE_CHECKING, Any, TypeVar, final
 
 from xdsl.utils.exceptions import VerifyException
 
@@ -444,30 +445,98 @@ def get_canonicalization_patterns(cls) -> tuple[RewritePattern, ...]:
         raise NotImplementedError()
 
 
+class MemoryEffectKind(Enum):
+    """
+    The kind of side effect an operation can have.
+
+    MLIR has a more detailed version of this, able to tie effects to specfic resources or
+    values. Here, everything has its effect on the universe.
+    """
+
+    READ = auto()
+    """
+    Indicates that the operation reads from some resource. A 'read' effect implies only
+    dereferencing of the resource, and not any visible mutation.
+    """
+
+    WRITE = auto()
+    """
+    Indicates that the operation writes to some resource. A 'write' effect implies only
+    mutating a resource, and not any visible dereference or read.
+    """
+
+    ALLOC = auto()
+    """
+    Indicates that the operation allocates from some resource. An 'allocate' effect
+    implies only allocation of the resource, and not any visible mutation or dereference.
+    """
+
+    FREE = auto()
+    """
+    Indicates that the operation frees some resource that has been allocated. A 'free'
+    effect implies only de-allocation of the resource, and not any visible allocation,
+    mutation or dereference.
+    """
+
+
 class MemoryEffect(OpTrait):
     """
     A trait that enables operations to expose their side-effects or absence thereof.
-
-    NB: The MLIR implementation further allows to describe what *kind* of side-effects
-    an operation has, e.g., read-only, or allocation.
-    This one is a stripped down version for now, just saying if there are any
-    side-effects or not.
     """
 
     @classmethod
     @abc.abstractmethod
-    def has_effects(cls, op: Operation) -> bool:
+    def get_effects(cls, op: Operation) -> set[MemoryEffectKind] | None:
+        """
+        Returns the concrete side effects of the operation.
+
+        Return None if the operation cannot conclude - interpreted as if the operation
+        had no MemoryEffect interface in the first place.
+        """
         raise NotImplementedError()
 
+    @final
+    @classmethod
+    def has_effects(cls, op: Operation) -> bool:
+        """
+        Returns if the operation has any side effects.
+        """
+        effects = cls.get_effects(op)
+        return (effects is not None) and len(effects) > 0
+
 
-def is_side_effect_free(op: Operation):
+def only_has_effect(op: Operation, effect: MemoryEffectKind) -> bool:
+    """
+    Returns if the operation has the given side effects and no others.
+    """
+    return get_effects(op) == {effect}
+
+
+def is_side_effect_free(op: Operation) -> bool:
     """
     Boilerplate helper to check if a generic operation is side effect free for sure.
     """
-    # If it doesn't say, safely assume it has side effects.
-    if not (trait := op.get_trait(MemoryEffect)):
-        return False
-    return not trait.has_effects(op)
+    return get_effects(op) == set()
+
+
+def get_effects(op: Operation) -> set[MemoryEffectKind] | None:
+    """
+    Helper to get known side effects of an operation, including recursive effects.
+    None means that the operation has unknown effects, for safety.
+    """
+
+    effect_interfaces = op.get_traits_of_type(MemoryEffect)
+    if not effect_interfaces:
+        return None
+
+    effects = set[MemoryEffectKind]()
+    for it in op.get_traits_of_type(MemoryEffect):
+        it_effects = it.get_effects(op)
+        if it_effects is None:
+            return None
+        effects.update(it_effects)
+
+    return effects
 
 
 class NoMemoryEffect(MemoryEffect):
@@ -476,28 +545,54 @@ class NoMemoryEffect(MemoryEffect):
     """
 
     @classmethod
-    def has_effects(cls, op: Operation) -> bool:
-        return False
+    def get_effects(cls, op: Operation) -> set[MemoryEffectKind]:
+        return set()
+
+
+class MemoryReadEffect(MemoryEffect):
+    """
+    A trait that signals that an operation always and only have read side effects.
+    """
+
+    @classmethod
+    def get_effects(cls, op: Operation) -> set[MemoryEffectKind]:
+        return {MemoryEffectKind.READ}
+
+
+class MemoryWriteEffect(MemoryEffect):
+    """
+    A trait that signals that an operation always and only have write side effects.
+    """
+
+    @classmethod
+    def get_effects(cls, op: Operation) -> set[MemoryEffectKind]:
+        return {MemoryEffectKind.WRITE}
 
 
 class RecursiveMemoryEffect(MemoryEffect):
     """
     A trait that signals that an operation has the side effects of its contained
     operations.
-
-    NB: Upstream, this a separate class, but in our current binary side effect
-    implementation, it's easier to have it this way in my opinion.
     """
 
     @classmethod
-    def has_effects(cls, op: Operation) -> bool:
-        if not op.regions:
-            return True
-        return not all(is_side_effect_free(o) for r in op.regions for o in r.walk())
+    def get_effects(cls, op: Operation):
+        effects = set[MemoryEffectKind]()
+        for r in op.regions:
+            for b in r.blocks:
+                for child_op in b.ops:
+                    child_effects = get_effects(child_op)
+                    if child_effects is None:
+                        return None
+                    effects.update(child_effects)
+        return effects
 
 
 class Pure(NoMemoryEffect):
-    """A trait that signals that an operation has no side effects."""
+    """
+    In MLIR, Pure is NoMemoryEffect + AlwaysSpeculatable, but the latter is nowhere to be
+    found here.
+    """
 
 
 class HasInsnRepresentation(OpTrait, abc.ABC):