Add an integer range analysis (#17)

math-fehr · web-flow · commit c579002c3690 · 2025-06-21T06:43:15.000+02:00
diff --git a/asl_xdsl/analysis/__init__.py b/asl_xdsl/analysis/__init__.py
diff --git a/asl_xdsl/analysis/integer_range.py b/asl_xdsl/analysis/integer_range.py
@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+from xdsl.ir import Operation, Region, SSAValue
+from xdsl.traits import OpTrait
+
+
+@dataclass(frozen=True)
+class IntegerRange:
+    """
+    The range of an integer value, defined by its lower and upper bounds (inclusive).
+    If either bound is None, it means that the bound is not known.
+    """
+
+    lower_bound: int | None
+    upper_bound: int | None
+
+    def is_empty(self) -> bool:
+        """Check if the range is empty."""
+        if self.lower_bound is None or self.upper_bound is None:
+            return False
+        return self.lower_bound > self.upper_bound
+
+    def get_as_constant(self) -> int | None:
+        """
+        Get the range as a constant integer value if it is a single value range.
+        If the range is not a single value, returns None.
+        """
+        if self.lower_bound is not None and self.upper_bound is not None:
+            if self.lower_bound == self.upper_bound:
+                return self.lower_bound
+        return None
+
+    @staticmethod
+    def top() -> IntegerRange:
+        """
+        Get the top range, which is the range that covers all possible integer values.
+        This is used when no specific range is known for a value.
+        """
+        return IntegerRange(None, None)
+
+    @staticmethod
+    def bottom() -> IntegerRange:
+        """
+        Get the bottom range, which is the empty range.
+        This is used when a value is known to be outside of any possible integer range.
+        """
+        return IntegerRange(1, -1)
+
+    def __contains__(self, value: int) -> bool:
+        """Check if a value is within the range."""
+        if self.lower_bound is not None:
+            if value < self.lower_bound:
+                return False
+        if self.upper_bound is not None:
+            if value > self.upper_bound:
+                return False
+        return True
+
+    def __or__(self, other: IntegerRange) -> IntegerRange:
+        """Combine two integer ranges into one that covers both ranges."""
+        if self.lower_bound is None:
+            lower_bound = other.lower_bound
+        elif other.lower_bound is None:
+            lower_bound = self.lower_bound
+        else:
+            lower_bound = min(self.lower_bound, other.lower_bound)
+
+        if self.upper_bound is None:
+            upper_bound = other.upper_bound
+        elif other.upper_bound is None:
+            upper_bound = self.upper_bound
+        else:
+            upper_bound = max(self.upper_bound, other.upper_bound)
+
+        return IntegerRange(lower_bound, upper_bound)
+
+    def __and__(self, other: IntegerRange) -> IntegerRange:
+        """Intersect two integer ranges into one that covers the intersection."""
+        if self.lower_bound is None:
+            lower_bound = other.lower_bound
+        elif other.lower_bound is None:
+            lower_bound = self.lower_bound
+        else:
+            lower_bound = max(self.lower_bound, other.lower_bound)
+
+        if self.upper_bound is None:
+            upper_bound = other.upper_bound
+        elif other.upper_bound is None:
+            upper_bound = self.upper_bound
+        else:
+            upper_bound = min(self.upper_bound, other.upper_bound)
+
+        return IntegerRange(lower_bound, upper_bound)
+
+
+@dataclass
+class IntegerRangeAnalysis:
+    """An analysis that contains the integer ranges of integer values."""
+
+    ranges: dict[SSAValue, IntegerRange] = field(
+        default_factory=dict[SSAValue, IntegerRange]
+    )
+
+    def get_range(self, value: SSAValue) -> IntegerRange:
+        """
+        Get the integer range of a value.
+        If the value is not an integer type, returns None.
+        """
+        return self.ranges.get(value, IntegerRange.top())
+
+    def set_range(self, value: SSAValue, integer_range: IntegerRange) -> None:
+        """
+        Set the integer range of a value.
+        If the value is not an integer type, this will raise an error.
+        """
+        # No need to set the top range, as it is the default
+        if integer_range == IntegerRange.top():
+            return
+        self.ranges[value] = integer_range
+
+    def compute_operation_analysis(self, op: Operation) -> None:
+        """
+        Compute the integer ranges of the operation's results.
+        This function will call the IntegerRangeTrait.compute_analysis method
+        if the operation has the IntegerRangeTrait trait.
+        """
+        trait = op.get_trait(IntegerRangeTrait)
+        if trait is None:
+            return
+
+        trait.compute_analysis(op, self)
+
+    def compute_single_block_region_analysis(self, region: Region) -> None:
+        """
+        Compute the integer ranges of all the SSA values in an SSACFG region
+        with a single block.
+        """
+        assert len(region.blocks) == 1, "Region must have a single block"
+        block = region.block
+
+        for op in block.ops:
+            self.compute_operation_analysis(op)
+
+
+class IntegerRangeTrait(OpTrait, ABC):
+    """
+    A trait that indicates that an operation can be used for an integer range analysis.
+    In practice, this means that we can compute the lower and upper bounds of the
+    operation's integer results based on its operands and attributes ranges.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def compute_analysis(op: Operation, analysis: IntegerRangeAnalysis):
+        """
+        Compute the integer ranges of the operation's results based on the
+        integer ranges of its operands and attributes.
+        """
+        ...
diff --git a/asl_xdsl/dialects/asl.py b/asl_xdsl/dialects/asl.py
@@ -45,6 +45,12 @@
     SymbolUserOpInterface,
 )
 
+from asl_xdsl.analysis.integer_range import (
+    IntegerRange,
+    IntegerRangeAnalysis,
+    IntegerRangeTrait,
+)
+
 
 @irdl_attr_definition
 class ConstraintExactAttr(ParametrizedAttribute):
@@ -308,6 +314,13 @@ def print_parameters(self, printer: Printer) -> None:
         printer.print(">")
 
 
+class ConstantIntIntegerRangeTrait(IntegerRangeTrait):
+    @staticmethod
+    def compute_analysis(op: Operation, analysis: IntegerRangeAnalysis) -> None:
+        assert isinstance(op, ConstantIntOp), "Expected ConstantIntOp"
+        analysis.set_range(op.res, IntegerRange(op.value.data, op.value.data))
+
+
 @irdl_op_definition
 class ConstantIntOp(IRDLOperation):
     """A constant arbitrary-sized integer operation."""
@@ -317,6 +330,8 @@ class ConstantIntOp(IRDLOperation):
     value = prop_def(builtin.IntAttr)
     res = result_def(IntegerType)
 
+    traits = traits_def(ConstantIntIntegerRangeTrait())
+
     def __init__(
         self, value: int | builtin.IntAttr, attr_dict: Mapping[str, Attribute] = {}
     ):
@@ -573,6 +588,21 @@ class AlignIntOp(BinaryIntOp):
     name = "asl.align_int"
 
 
+class MoxPow2IntIntegerRangeTrait(IntegerRangeTrait):
+    @staticmethod
+    def compute_analysis(op: Operation, analysis: IntegerRangeAnalysis) -> None:
+        assert isinstance(op, ModPow2IntOp), "Expected ModPow2IntOp"
+        rhs_upper = analysis.get_range(op.rhs).upper_bound
+        lhs_bounds = analysis.get_range(op.lhs)
+        if rhs_upper is None:
+            # If the rhs is unbounded, the result is not more bounded than the lhs.
+            analysis.set_range(op.res, lhs_bounds)
+            return
+
+        bound_from_rhs = IntegerRange(0, 2**rhs_upper - 1)
+        analysis.set_range(op.res, lhs_bounds & bound_from_rhs)
+
+
 @irdl_op_definition
 class ModPow2IntOp(BinaryIntOp):
     """
@@ -582,6 +612,8 @@ class ModPow2IntOp(BinaryIntOp):
 
     name = "asl.mod_pow2_int"
 
+    traits = traits_def(MoxPow2IntIntegerRangeTrait())
+
 
 @irdl_op_definition
 class IsPow2IntOp(IRDLOperation):
diff --git a/asl_xdsl/tools/asl_opt.py b/asl_xdsl/tools/asl_opt.py
@@ -17,6 +17,16 @@ def register_all_dialects(self):
         self.ctx.load_dialect(ASLDepDialect)
 
     def register_all_passes(self):
+        def get_integer_range_analysis_pass():
+            from asl_xdsl.transforms.test_integer_range_analysis import (
+                TestIntegerRangeAnalysis,
+            )
+
+            return TestIntegerRangeAnalysis
+
+        self.register_pass(
+            "test-integer-range-analysis", get_integer_range_analysis_pass
+        )
         return super().register_all_passes()
 
     def register_all_targets(self):
diff --git a/asl_xdsl/transforms/__init__.py b/asl_xdsl/transforms/__init__.py
diff --git a/asl_xdsl/transforms/test_integer_range_analysis.py b/asl_xdsl/transforms/test_integer_range_analysis.py
@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+
+from xdsl.context import Context
+from xdsl.dialects.builtin import ArrayAttr, IntAttr, ModuleOp, NoneAttr
+from xdsl.ir import Attribute
+from xdsl.passes import ModulePass
+
+from asl_xdsl.analysis.integer_range import IntegerRange, IntegerRangeAnalysis
+
+
+def bound_to_attr(bound: int | None) -> Attribute:
+    """Convert a bound (int or None) to an attribute."""
+    if bound is None:
+        return NoneAttr()
+    return IntAttr(bound)
+
+
+def range_to_attr(integer_range: IntegerRange) -> Attribute:
+    """Convert an IntegerRange to an attribute."""
+    return ArrayAttr(
+        [
+            bound_to_attr(integer_range.lower_bound),
+            bound_to_attr(integer_range.upper_bound),
+        ]
+    )
+
+
+@dataclass(frozen=True)
+class TestIntegerRangeAnalysis(ModulePass):
+    """
+    Test the integer range analysis pass by adding integer range analysis information
+    to the module using attributes.
+    """
+
+    name = "test-integer-range-analysis"
+
+    def apply(self, ctx: Context, op: ModuleOp) -> None:
+        analysis = IntegerRangeAnalysis()
+        analysis.compute_single_block_region_analysis(op.body)
+
+        for sub_op in op.walk():
+            if not sub_op.results:
+                continue
+            result_ranges = [analysis.get_range(result) for result in sub_op.results]
+            result_ranges_attrs = ArrayAttr(
+                [range_to_attr(result) for result in result_ranges]
+            )
+            sub_op.attributes["__integer_ranges"] = result_ranges_attrs
diff --git a/tests/filecheck/analysis/test-integer-range-analysis.mlir b/tests/filecheck/analysis/test-integer-range-analysis.mlir
@@ -0,0 +1,13 @@
+// RUN: asl-opt %s -p=test-integer-range-analysis | filecheck %s
+
+builtin.module {
+
+    // CHECK: %unknown = "test.op"() {__integer_ranges = {{[}}[#none, #none]]} : () -> !asl.int
+    %unknown = "test.op"() : () -> !asl.int
+    
+    // CHECK-NEXT: %cst16 = asl.constant_int 16 {__integer_ranges = {{[}}[#builtin.int<16>, #builtin.int<16>]]}
+    %cst16 = asl.constant_int 16
+    
+    // CHECK-NEXT: %res = asl.mod_pow2_int %unknown, %cst16 : (!asl.int, !asl.int) -> !asl.int {__integer_ranges = {{[}}[#builtin.int<0>, #builtin.int<65535>]]}
+    %res = asl.mod_pow2_int %unknown, %cst16 : (!asl.int, !asl.int) -> !asl.int
+}
