[Unity][IR] Purity Tracking

include/tvm/relax/analysis.h

@@ -438,6 +438,20 @@ TVM_DLL relay::OpPatternKind AnalyzeOpPatternKind(const tir::PrimFunc& func);
  */
 TVM_DLL bool HasReshapePattern(const tir::PrimFunc& func);
 
+/*!
+ * \brief Check if the given expression (likely a function body) contains any impure calls.
+ * \param expr The expression to be examined. If expr is a function, we check the body.
+ * \param own_name (Optional.) If we are checking a recursive function body,
+ *   the caller can pass the function's name so recursive calls
+ *   can be ignored in the check (must be a Var or GlobalVar).
+ * \return A boolean indicating if the expression contains any impure calls.
+ * \note Relies on StructInfo annotations, so ensure that the module has been normalized first.
+ *   Also, an impure call in a *nested* function does *not* mean that the outer expression contains
+ *   an impure call--it only does if the nested function is *later called*.
+ */
+TVM_DLL bool ContainsImpureCall(const Expr& expr,
+                                const Optional<Expr>& own_name = Optional<Expr>(nullptr));
+
 /*!
  * \brief Check if the IRModule is well formed.
  *

include/tvm/relax/expr.h

@@ -920,10 +920,13 @@ class FunctionNode : public BaseFuncNode {
   Expr body;
   /*! \brief The return type of the function. */
   StructInfo ret_struct_info;
+  /*! \brief Whether the function is annotated as pure or not. */
+  bool is_pure;
 
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("params", &params);
     v->Visit("body", &body);
+    v->Visit("is_pure", &is_pure);
     v->Visit("ret_struct_info", &ret_struct_info);
     v->Visit("attrs", &attrs);
     v->Visit("struct_info_", &struct_info_);
@@ -934,15 +937,16 @@ class FunctionNode : public BaseFuncNode {
   bool SEqualReduce(const FunctionNode* other, SEqualReducer equal) const {
     equal->MarkGraphNode();
     return equal.DefEqual(params, other->params) && equal(body, other->body) &&
-           equal(ret_struct_info, other->ret_struct_info) && equal(attrs, other->attrs) &&
-           equal(struct_info_, other->struct_info_);
+           equal(ret_struct_info, other->ret_struct_info) && equal(is_pure, other->is_pure) &&
+           equal(attrs, other->attrs) && equal(struct_info_, other->struct_info_);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce->MarkGraphNode();
     hash_reduce.DefHash(params);
     hash_reduce(body);
     hash_reduce(ret_struct_info);
+    hash_reduce(is_pure);
     hash_reduce(attrs);
     hash_reduce(struct_info_);
   }
@@ -956,14 +960,16 @@ class FunctionNode : public BaseFuncNode {
 class Function : public BaseFunc {
  public:
   TVM_DLL explicit Function(Array<Var> params, Expr body, Optional<StructInfo> ret_struct_info,
-                            DictAttrs attrs = NullValue<DictAttrs>(), Span span = Span());
+                            bool is_pure = true, DictAttrs attrs = NullValue<DictAttrs>(),
+                            Span span = Span());
 
   /*!
    * \brief Mimics the constructor but without body Expr.
-   * \note ret_struct_info is required, since it can not deduced by the body
+   * \note ret_struct_info is required, since it can not deduced by the body.
    */
   TVM_DLL static Function CreateEmpty(Array<Var> params, StructInfo ret_struct_info,
-                                      DictAttrs attrs = NullValue<DictAttrs>(), Span span = Span());
+                                      bool is_pure = true, DictAttrs attrs = NullValue<DictAttrs>(),
+                                      Span span = Span());
 
   TVM_DEFINE_OBJECT_REF_METHODS(Function, BaseFunc, FunctionNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionNode);
@@ -985,6 +991,12 @@ constexpr const char* kComposite = "Composite";
 constexpr const char* kPartitionedFromPattern = "PartitionedFromPattern";
 /*! \brief The required workspace for an external function. */
 constexpr const char* kWorkspaceSize = "WorkspaceSize";
+
+// Note: in the future, we prefer snake_case instead of CamelCase for attributes.
+// Past ones will be kept for backwards compatibility.
+/*! \brief Override checking purity for this function and treat as pure
+ * (is_pure must be set to true) */
+constexpr const char* kForcePure = "relax.force_pure";
 }  // namespace attr
 
 /*! \brief The extern function, which can represent packed function. */

include/tvm/relax/struct_info.h

@@ -296,6 +296,12 @@ class FuncStructInfoNode : public StructInfoNode {
    *       ret should be ObjectStructInfo()
    */
   Optional<StructInfoDeriveFunc> derive_func;
+  /*!
+   * \brief Whether the function is pure.
+   * \note This parameter should be set to true only if the function is pure on all inputs.
+   *   If the function _may_ have visible side effects, set it to false.
+   */
+  bool purity;
 
   /*!
    * \return Whether the func struct info is opaque.
@@ -308,16 +314,18 @@ class FuncStructInfoNode : public StructInfoNode {
     v->Visit("ret", &ret);
     v->Visit("derive_func", &derive_func);
     v->Visit("span", &span);
+    v->Visit("purity", &purity);
   }
 
   bool SEqualReduce(const FuncStructInfoNode* other, SEqualReducer equal) const {
     return equal.DefEqual(params, other->params) && equal(ret, other->ret) &&
-           equal(derive_func, other->derive_func);
+           equal(purity, other->purity) && equal(derive_func, other->derive_func);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce.DefHash(params);
     hash_reduce(ret);
+    hash_reduce(purity);
     hash_reduce(derive_func);
   }
 
@@ -335,34 +343,42 @@ class FuncStructInfo : public StructInfo {
    * \brief Constructor from parameter struct info and return value struct info.
    * \param params The struct info of function parameters.
    * \param ret The return value struct info.
+   * \param purity The purity of the function (true by default).
    * \param span The span of the AST.
    *
    * \note If the ret contains variables(tir::Var and relax::Var), they must be deducible from
    * params. If you are unsure, you can always erase ret to static.
    */
-  TVM_DLL FuncStructInfo(Array<StructInfo> params, StructInfo ret, Span span = Span());
+  TVM_DLL FuncStructInfo(Array<StructInfo> params, StructInfo ret, bool purity = true,
+                         Span span = Span());
 
   /*!
    * \brief Constructing an opaque function struct info using derive_func.
    *
    * \param derive_func Derivation function.
+   * \param purity The purity of the function
+   *   (false by default: most external functions are not pure).
    * \param span The span of the AST.
    *
    * \return The FuncStructInfo for opaque packedfunc.
    * \note Defaults to an derive func that always return ObjectStructInfo if not specified.
    */
-  TVM_DLL static FuncStructInfo OpaqueFunc(StructInfoDeriveFunc derive_func, Span span = Span());
+  TVM_DLL static FuncStructInfo OpaqueFunc(StructInfoDeriveFunc derive_func, bool purity = false,
+                                           Span span = Span());
 
   /*!
    * \brief Construct an opaque function using from return struct info.
    *
    * \param ret The struct info of the return value.
+   * \param purity The purity of the function
+   *   (false by default: most external functions are not pure).
    * \param span The span of the AST.
    *
    * \return The FuncStructInfo for opaque packedfunc.
    * \note Defaults to an derive func that always return ObjectStructInfo if not specified.
    */
-  TVM_DLL static FuncStructInfo OpaqueFunc(StructInfo ret = ObjectStructInfo(), Span span = Span());
+  TVM_DLL static FuncStructInfo OpaqueFunc(StructInfo ret = ObjectStructInfo(), bool purity = false,
+                                           Span span = Span());
 
   TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FuncStructInfo, StructInfo, FuncStructInfoNode);
 };

include/tvm/relax/transform.h

@@ -83,6 +83,19 @@ TVM_DLL Pass LambdaLift();
  */
 TVM_DLL Pass ToNonDataflow();
 
+/*!
+ * \brief Activate force_pure on all pure functions in the module
+ * and unwrap all pure override ops into the normal versions.
+ *
+ * This effectively means that there will be no more purity tracking,
+ * useful for low-level code generation.
+ *
+ * \return The Pass.
+ *
+ * \note Should be used after ToNonDataflow()
+ */
+TVM_DLL Pass RemovePurityChecking();
+
 /*!
  * \brief Perform explicit tensor allocation for call_tir and call_dps_packed.
  *

include/tvm/relax/utils.h

@@ -81,6 +81,18 @@ TVM_DLL bool IsBoolStructInfo(const StructInfo& sinfo, bool permit_unknown_rank
  */
 TVM_DLL bool IsLeafOrTuple(const Expr& expr);
 
+/*!
+ * \brief Check if the given Call node is an impure operation. If the callee is a general
+ * expression, this simply requires checking the purity field of the FuncStructInfo. If it is an Op,
+ * then this checks the `fPurity` field.
+ *
+ * \param call The input call
+ *
+ * \return True iff the call is impure (definitely or possibly results in a visible side effect).
+ *   That is, a call is considered pure only if definitely does not result in a visible side effect.
+ */
+TVM_DLL bool IsImpureCall(const Call& call);
+
 /*!
  * \brief Copy the given function. All variables that are bound inside the original function
  *  would be copied to satisfy the restriction in the well-formed check: Variables in

include/tvm/script/ir_builder/relax/frame.h

@@ -97,7 +97,8 @@ class FunctionFrameNode : public SeqExprFrameNode {
    *       take the specified `ret_struct_info`.
    */
   Optional<tvm::relax::StructInfo> ret_struct_info;
-
+  /*! \brief Whether the function is annotated as pure */
+  Optional<Bool> is_pure;
   /*! \brief The function attributes. */
   Map<String, ObjectRef> attrs;
   /*! \brief The block builder to create Relax function. */
@@ -108,6 +109,7 @@ class FunctionFrameNode : public SeqExprFrameNode {
     v->Visit("name", &name);
     v->Visit("params", &params);
     v->Visit("ret_struct_info", &ret_struct_info);
+    v->Visit("is_pure", &is_pure);
     v->Visit("attrs", &attrs);
     v->Visit("binding_blocks", &binding_blocks);
     v->Visit("output", &output);

include/tvm/script/ir_builder/relax/ir.h

@@ -57,6 +57,12 @@ TVM_DLL void FuncName(const String& name);
  */
 TVM_DLL void FuncAttrs(Map<String, ObjectRef> attrs);
 
+/*!
+ * \brief Specify the purity of the last function frame.
+ * \param purity Whether the function is pure.
+ */
+TVM_DLL void FuncIsPure(bool purity);
+
 /*!
  * \brief Specify the return struct info of the last function frame.
  * \param ret_sinfo The return struct info.

python/tvm/relax/__init__.py

@@ -56,7 +56,7 @@
 from .exec_builder import ExecBuilder
 
 # Operator
-from .op.base import call_tir, call_dps_packed
+from .op.base import call_tir, call_pure_packed, call_dps_packed
 
 # BlockBuilder
 from .block_builder import BlockBuilder

python/tvm/relax/analysis/analysis.py

@@ -21,7 +21,7 @@
 configuring the passes and scripting them in Python.
 """
 
-from typing import Dict, List, Union, Callable
+from typing import Dict, List, Optional, Union, Callable
 from enum import IntEnum
 
 import tvm
@@ -327,6 +327,34 @@ def has_reshape_pattern(func: tir.PrimFunc) -> bool:
     return _ffi_api.has_reshape_pattern(func)  # type: ignore
 
 
+def contains_impure_call(expr: Expr, own_name: Optional[Union[Var, GlobalVar]] = None) -> bool:
+    """
+    Check if the given expression (likely a function body) contains any impure calls.
+
+    Parameter
+    ---------
+    expr : Expr
+        The expression to be examined. If expr is a function, we check the body.
+
+    own_name : Var or GlobalVar (optional)
+        For a recursive function, the analysis can ignore the self-calls
+        for checking purity.
+
+    Returns
+    -------
+    ret : bool
+        True if there is an impure call
+        (call to a function that may have visible side effects).
+
+    Notes
+    -----
+    Relies on StructInfo annotations, so ensure that the module has been normalized first.
+    Also, an impure call in a *nested* function does *not* mean that the outer expression contains
+    an impure call--it only does if the nested function is *later called*.
+    """
+    return _ffi_api.contains_impure_call(expr, own_name)
+
+
 def get_var2val(func: Function) -> Dict[Var, Expr]:
     """
     Get a mapping from Var to Expr for each variable in the function.

python/tvm/relax/backend/contrib/cutlass.py

@@ -385,7 +385,7 @@ def __init__(self, mod):
     def visit_function_(self, f):
         if f.attrs is None or "Composite" not in f.attrs:
             body = super().visit_expr(f.body)
-            new_f = Function(f.params, body, f.ret_struct_info, f.attrs, f.span)
+            new_f = Function(f.params, body, f.ret_struct_info, f.is_pure, f.attrs, f.span)
 
             if f.attrs and "global_symbol" in f.attrs and "cutlass" in f.attrs["global_symbol"]:
                 composite_func = body.blocks[0].bindings[0].value

python/tvm/relax/expr.py

@@ -560,29 +560,40 @@ class Function(BaseFunc, Scriptable):
     params: List[Var]
     body: Expr
     ret_struct_info: StructInfo
+    is_pure: bool
     attrs: Optional[tvm.ir.DictAttrs]
 
     def __init__(
         self,
         params: List[Var],
         body: Expr,
         ret_struct_info: Optional[StructInfo] = None,
+        is_pure: Optional[bool] = True,
         attrs: Optional[tvm.ir.DictAttrs] = None,
         span: Optional[Span] = None,
     ) -> None:
         self.__init_handle_by_constructor__(
-            _ffi_api.Function, params, body, ret_struct_info, attrs, span  # type: ignore
-        )
+            _ffi_api.Function,
+            params,
+            body,
+            ret_struct_info,
+            is_pure,
+            attrs,
+            span,  # type: ignore
+        )  # type: ignore
 
     @staticmethod
     def create_empty(
         params: List[Var],
         ret_struct_info: StructInfo,
+        is_pure: Optional[bool] = True,
         attrs: Optional[tvm.ir.DictAttrs] = None,
         span: Optional[Span] = None,
     ):
         """Construct a relax.Function but without body"""
-        return _ffi_api.FunctionCreateEmpty(params, ret_struct_info, attrs, span)  # type: ignore
+        return _ffi_api.FunctionCreateEmpty(
+            params, ret_struct_info, is_pure, attrs, span
+        )  # type: ignore
 
     def __call__(self, *args):
         """Invoke the global function.