[Unity][Analysis] Analysis for detecting recursion in Relax (#14149)

* DFS based attempt to detect mutual recursion

* Use Johnson's circuit-detecting algorithm instead

* Fix control flow test

* Detect all recursion anyway

* Add new test cases for simple recursion

* Fix mistake in test case

* Include missing dependencies

* Remove trailing whitespace

* Dependencies are simply references, not necessarily calls

* More trailing whitespace

* Newline at end of file

* Fix spacing in docstring

Co-authored-by: Siyuan Feng <[email protected]>

---------

Co-authored-by: Siyuan Feng <[email protected]>

include/tvm/relax/analysis.h

@@ -308,6 +308,31 @@ TVM_DLL tvm::Array<Var> AllVars(const Expr& expr);
  */
 TVM_DLL tvm::Array<GlobalVar> AllGlobalVars(const Expr& expr);
 
+/*!
+ * \brief Find all sets of recursive or mutually recursive functions in the module.
+ *
+ * Two or more functions are mutually recursive if there is some cycle of references
+ * among them. For example, if there are two functions A and B, they are
+ * mutually recursive if A calls B and B calls A. Another case would be with
+ * three functions A, B, and C, where A calls B, B calls C, and C calls A.
+ *
+ * (Note that functions do not have to call each other to reference each other.
+ * For example, if a function returns another function, that is still a reference
+ * that could potentially be recursive, even without a call.)
+ *
+ * If a function is simply recursive and not mutually recursive with any other,
+ * it will be reported as a group by itself.
+ *
+ * \param m The module
+ *
+ * \return List of all groups of mutually recursive functions.
+ *     Each member of the result is a list of functions in the module
+ *     that are all mutually recursive.
+ *     If a function is simply recursive and not mutually recursive with any other,
+ *     then it will be listed as a group by itself.
+ */
+TVM_DLL tvm::Array<tvm::Array<GlobalVar>> DetectRecursion(const IRModule& m);
+
 /*!
  * \brief Analyze var -> value mapping from VarBindings.
  *

python/tvm/relax/analysis/analysis.py

@@ -24,11 +24,12 @@
 from typing import Dict, List, Union, Callable
 from enum import IntEnum
 
+import tvm
 from tvm import tir
 from tvm import IRModule
 from tvm.relax.ty import Type
 from tvm.relax.struct_info import StructInfo, FuncStructInfo
-from tvm.relax.expr import DataflowBlock, Var, Expr, Function, Call, Binding
+from tvm.relax.expr import DataflowBlock, Var, GlobalVar, Expr, Function, Call, Binding
 from tvm.tir import IndexMap, PrimFunc, Block, Buffer
 from . import _ffi_api
 
@@ -319,3 +320,35 @@ def suggest_layout_transforms(
         assert isinstance(transform, IndexMap)
         write_buffer_index_maps.append(transform)
     return _ffi_api.suggest_layout_transforms(func, write_buffer_index_maps)  # type: ignore
+
+
+def detect_recursion(mod: tvm.IRModule) -> List[List[GlobalVar]]:
+    """
+    Find all sets of recursive or mutually recursive functions in the module.
+
+    Two or more functions are mutually recursive if there is some cycle of references
+    among them. For example, if there are two functions A and B, they are
+    mutually recursive if A calls B and B calls A. Another case would be with
+    three functions A, B, and C, where A calls B, B calls C, and C calls A.
+
+    (Note that functions do not have to call each other to reference each other.
+    For example, if a function returns another function, that is still a reference
+    that could potentially be recursive, even without a call.)
+
+
+    If a function is simply recursive and not mutually recursive with any other,
+    it will be reported as a group by itself.
+
+    Parameters
+    ----------
+    mod: The module
+
+    Returns
+    -------
+    ret: List[List[GlobalVar]]
+        Each member of the list is a list of global functions
+        that references each other mutually recursively.
+        If a function is simply recursive and not mutually recursive
+        with any other, it will be a singleton in this list.
+    """
+    return _ffi_api.detect_recursion(mod)  # type: ignore