[IR] Improve pass infra

onnxscript/ir/passes/_pass_infra.py

@@ -66,16 +66,47 @@
 
 
 class PassBase(abc.ABC):
-    """Base class for all passes.
+    """Base class for all passes."""
 
-    Class attributes:
-        in_place: Whether the pass modifies the model in place.
-    """
+    @property
+    def in_place(self) -> bool:
+        """Whether the pass modifies the model in place."""
+        return True
 
-    in_place: bool = True
+    @property
+    def destructive(self) -> bool:
+        """Whether the pass will destroy the input model when ``in_place=False``."""
+        return False
 
     def __call__(self, model: ir.Model) -> PassResult:
-        return self.call(model)
+        # Check preconditions
+        try:
+            self.requires(model)
+        except PreconditionError:
+            raise
+        except Exception as e:
+            raise PreconditionError(
+                f"Pre-condition for pass '{self.__class__.__name__}' failed"
+            ) from e
+
+        result = self.call(model)
+
+        # Check postconditions
+        try:
+            self.ensures(model)
+        except PostconditionError:
+            raise
+        except Exception as e:
+            raise PostconditionError(
+                f"Post-condition for pass '{self.__class__.__name__}' failed"
+            ) from e
+
+        if not isinstance(result, PassResult):
+            raise TypeError(
+                f"The result of the pass '{self.__class__.__name__}' should be type PassResult. "
+                "Please create one with ir.passes.PassResult()."
+            )
+        return result
 
     @abc.abstractmethod
     def call(self, model: ir.Model) -> PassResult:
@@ -97,29 +128,37 @@
         del model  # Unused
 
 
-class PassManager:
+class PassManager(PassBase):
     """Pass manager for the IR.
 
-    The PassManager is a callable that runs a sequence of passes on a model.
+    The PassManager is a Pass that runs a sequence of passes on a model.
 
     Attributes:
         passes: The passes to run.
-        check_invariants: Whether to check invariants before and after each pass.
         steps: The number of times to run the passes.
     """
 
     def __init__(
         self,
         passes: Sequence[PassBase],
-        check_invariants: bool = False,
         steps: int = 1,
     ):
         # TODO(justinchuby): Implement constraints
         self.passes = list(passes)
-        self.check_invariants = check_invariants
         self.steps = steps
 
-    def __call__(self, model: ir.Model) -> PassResult:
+    @property
+    def in_place(self) -> bool:
+        """Whether the pass modifies the model in place."""
+        return all(pass_.in_place for pass_ in self.passes)
+
+    @property
+    def destructive(self) -> bool:
+        """Whether the pass will destroy the input model when ``in_place=False``."""
+        # This logic is a little conservative, but it is ok for now
+        return any(pass_.destructive for pass_ in self.passes)
+
+    def call(self, model: ir.Model) -> PassResult:
         """Run the set of passes `steps` number of times or until the graph stops changing."""
         overall_modified = False
         for step in range(self.steps):
@@ -137,36 +176,18 @@
         modified = False
         for i, pass_ in enumerate(self.passes):
             logger.debug("Running the %s-th pass '%s', (step %s)", i, pass_, step)
-
-            # 1. Check preconditions
-            if self.check_invariants:
-                try:
-                    pass_.requires(model)
-                except Exception as e:
-                    raise PreconditionError(f"Pre-condition failed for {pass_}") from e
-
-            # 2. Run the pass
             try:
                 pass_result = pass_(model)
+            except (PreconditionError, PostconditionError):
+                raise
             except Exception as e:
                 prev_pass_names = [str(p) for p in self.passes[:i]]
                 raise PassError(
                     f"An error occurred when running the '{pass_}' pass after the "
                     f"following passes: {prev_pass_names} during step {step}"
                 ) from e
-            if not isinstance(pass_result, PassResult):
-                raise TypeError(
-                    f"The result of the pass {pass_} should be type PassResult."
-                    "Please create one with ir.passes.PassResult()."
-                )
 
             model = pass_result.model
             modified = modified or pass_result.modified
 
-            # 3. Check postconditions
-            if self.check_invariants:
-                try:
-                    pass_.ensures(model)
-                except Exception as e:
-                    raise PostconditionError(f"Post-condition failed for {pass_}") from e
         return PassResult(model, modified)