[checkpoint] clarrify lowering does not support flttened result tuple types

[checkpoint] Don't allow FromExprInContext to be ffi-able.

Update tests/python/relay/test_backend_interpreter.py

Co-authored-by: Altan Haan <altanh@cs.washington.edu>

Update tests/python/relay/test_backend_interpreter.py

Co-authored-by: Altan Haan <altanh@cs.washington.edu>

Update tests/python/relay/test_backend_interpreter.py

Co-authored-by: Altan Haan <altanh@cs.washington.edu>

Update tests/python/relay/test_backend_interpreter.py

Co-authored-by: Altan Haan <altanh@cs.washington.edu>

[checkpoint] Fix target bug.

[checkpoint] Explicit hash struct

[checkpoint] flubbed add

[checkpoint] Shape functions always run on cpu

- Support target->lowered irmodule map in interpreter so  can keep device vs cpu ir functions separated.
- Cleanup flattening / construction
  of ADT/NDArray to mirror what we do for tuple types.
- Make sure unit tests compile clean for no signed comparison.

[checkpoint] logging to debug cuda failure

[checkpoint] No signed/unsigned compare.

Why is -Werror=sign-compare not part of the cmake setup???

[checkpoint] Include lint.

[checkpoint] Cleanup tuple flattening.

[checkpoint] First batch of Jared's comments, mostly fixup FromExpr.

[checkpoint] Fixup target resolution.

Somewhere early on I got confused about how the device context map works. Fixed.

[checkpoint] I will run black before every commit.
I will run black before every commit.

[checkpoint] Bug fixes.

[checkpoint] bug fixes

[checkpoint] can't pass my own tests!

[checkpoint] format

[checkpoint] typos

[checkpoint] bit the bullet and bring create_executor & evaluate together as much as possible

[checkpoint] doc lints

[checkpoint] sigh

[checkpoint] more python lints, hoist executor for adt tests to exploit cache

[checkpoint] Python lints

[checkpoint] lint fixes

[checkpoint] cleanup, nuke vlog stuff, collape create_executor() & evaluate()

[checkpoint] Bug fixes. AOT is broken.

[checkpoint] Separate eval-to-clousure and apply-closure phases at last

[checkpoint] Fix GetType recursion, get debug going.

[checkpoint] Audit python to collapse create_executor and evaluate phases

Just a few places where this doesn't work, seem harmless.

[checkpoint] Get interpreter working using tec::LowerTE, but no dynamic shapes.

 - Hide TECompiler impl inside te_compiler.cc. However I think it is already exposed
   into Python land so this probably won't be possible now.
 - Move 'optimize' pre-transforms from interpreter.py to interpreter.cc so can be
   applied uniformly to both mod and expr.
 - Don't push the expr into the mod in interpreter.py since it's done again in
   interpreter.cc. Instead just build the Call node with the reflected args.
 - Both the mod and the expr are prepared identically (same transforms, of which
   LowerTensorExpr should be one).
 - LowerTensorExpr can look through let-bound and global vars, eg
     let f = fn (..., Primitive=1) { ... } ... f(...)
     ==> @lowered_f = ... @lowered_f(...)
 - Lots of DLOGs that need to be removed or reorganized.

[checkpoint] Support shape functions.

TODO:
 - Unit tests.
 - Cleanup logging (VLOG?)
 - Don't build all prims on each apply.

[checkpoint] typo

[checkpoint] Don't allow evaling expr independently of preparing module.

TODO:
 - Make eval(mod, expr) the interface.
 - GlobalVar's don't line up.
 - Rework use of interpreter in fold_constant.cc to make clear
   it is evaling prim calls which have already been prepared.
 - Find a dynamic shape example that works at HEAD.
 - Unit tests.

[checkpoint] Interpreting expression with refs to module defs working

Commit to interpreter evaling expr w.r.t. mod in single phase. Thankfully
turns out no existing uses broke that assumption so we dodged a bullet.

Binding of expr into mod to be evaled is a mess, needs to be fixed.

Still can't confirm dynamic shapes working since don't have an
example working at HEAD.

Change to partial_eval needs to be tested, but smells ok.

[checkpoint] Dynamic shapes working

The use of TIRCallAttrs is pretty hacky but the shape function
calls are at least working.

Next is to tackle the 'build everything just to project one prim fun' problem.

[checkpoint] Cache built prims, make sure build with minimal deps.

[checkpoint] Cleanup expr-to-module hackery.

Author: mbs-octoml

docs/langref/relay_pattern.rst

@@ -406,7 +406,7 @@ Either match the first pattern or the second pattern.
 Domination
 **********
 
-Match child pattern, find a match for the parent pattern, insuring that the child ultimately dominates the parrent (i.e., no nodes outside the pattern use outputs of the parent), and that ever node betwen the child and the pattern matches the path pattern.
+Match child pattern, find a match for the parent pattern, insuring that the child ultimately dominates the parent (i.e., no nodes outside the pattern use outputs of the parent), and that ever node between the child and the pattern matches the path pattern.
 
 Function Pattern
 ****************

include/tvm/ir/module.h

@@ -36,6 +36,7 @@
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 namespace tvm {
@@ -253,6 +254,14 @@ class IRModuleNode : public Object {
   /*! \brief Helper function for registering a typedef's constructors */
   void RegisterConstructors(const GlobalTypeVar& var, const TypeData& type);
 
+  /*!
+   * \brief Returns a version of \p name which is unique amongst all function definitions in module.
+   *
+   * \param name The original name.
+   * \return Updated name which is unique.
+   */
+  String GetUniqueName(const String& name);
+
   /*! \brief A map from string names to global variables that
    * ensures global uniqueness.
    */
@@ -307,16 +316,38 @@ class IRModule : public ObjectRef {
   }
 
   /*!
-   * \brief Construct a module from a standalone expression.
+   * \brief Constructs a module from a standalone expression \p expr.
+   *
+   * If \p expr is a function it will be bound directly. Otherwise a function over the free
+   * variables of \p expr (possibly none) with \p expr as body is created and bound.
+   *
+   * The function is bound to, in preference order:
+   *  - The "global_symbol" attribute of \p expr, if it is a function with that attribute.
+   *  - 'main'
+   *  - A unique name derived from 'main' if 'main' is already bound in \p global_funcs.
    *
-   * Allows one to optionally pass a global function map and
-   * map of type definitions as well.
+   * Additional global functions and type definitions may be included in the result module.
+   *
+   * See also \p FromExpr.
    *
    * \param expr The expression to set as the main function to the module.
-   * \param global_funcs The global function map.
-   * \param type_definitions Map of global type definitions
+   * \param global_funcs The global function map. Default empty.
+   * \param type_definitions The global type definition map. Default empty.
+   * \param import_set Set of external modules already imported. Default empty.
+   *
+   * \returns A module with \p expr set as the main function, and the global var to which
+   * \p expr was bound (typcially 'main').
    *
-   * \returns A module with expr set as the main function.
+   * TODO(mbs): Does import_set and the bound global var need to be exposed via ffi?
+   */
+  static std::pair<IRModule, GlobalVar> FromExprInContext(
+      const RelayExpr& expr, const Map<GlobalVar, BaseFunc>& global_funcs = {},
+      const Map<GlobalTypeVar, TypeData>& type_definitions = {},
+      std::unordered_set<String> import_set = {});
+
+  /*!
+   * \brief As for \p FromExprInContext, but assuming \p expr is bound to 'main' and no
+   * imports.
    */
   TVM_DLL static IRModule FromExpr(const RelayExpr& expr,
                                    const Map<GlobalVar, BaseFunc>& global_funcs = {},

include/tvm/relay/interpreter.h

@@ -40,31 +40,11 @@
 #include <tvm/runtime/object.h>
 #include <tvm/target/target.h>
 
+#include <unordered_set>
+
 namespace tvm {
 namespace relay {
 
-/*!
- *\brief Create a Interpreter function that can
- *  evaluate an expression and produce a value.
- *
- * The resulting value can be passed to Python, making it easy to use
- * for testing and debugging.
- *
- * The interpreter interprets the program fragments not supported by the
- * TVM runtime, although the interpreter is naively implemented it uses
- * TVM operators for evaluating all operators.
- *
- * Our intent is that this will never be the most efficient implementation of
- * Relay's semantics, but a readable and clear one.
- *
- * \param mod The function module.
- * \param device The primary device that the interepreter runs on.
- * \param target Compiler target flag to compile the functions on the context.
- * \return A function that takes in an expression and returns a value.
- */
-runtime::TypedPackedFunc<ObjectRef(Expr)> CreateInterpreter(IRModule mod, Device device,
-                                                            Target target);
-
 /*! \brief The container type of Closures used by the interpreter. */
 class InterpreterClosureObj : public runtime::ClosureObj {
  public:
@@ -164,6 +144,52 @@ class ConstructorValue : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_METHODS(ConstructorValue, ObjectRef, ConstructorValueObj);
 };
 
+/*!
+ * \brief Returns a packed function over Relay expressions which will evaluate \p expr
+ * applied to those arguments, where \p expr is w.r.t. the definitions in \p mod.
+ *
+ * This function is intended to support the Python 'debug' executor.
+ *
+ * The given \p expr should have function type. The given \p mod may be empty or
+ * undefined if \p expr is self-contained. Relay arguments passed to the result
+ * packed function must be constants, references, or constructors/tuples over such.
+ * As much work as possible is done while constructing the result packed function, and
+ * that function may be reasonably efficiently applied multiple times without redoing
+ * unnecessary work.
+ *
+ * Primitives are lowered and compiled to packed functions for execution on \p device
+ * with properties given by \p target. All other Relay constructs are interpreted.
+ *
+ * The interpreter is intended to be a 'reference' implementation of the Relay semantics
+ * for testing and interactive use. It is not intended to be particularly efficient.
+ *
+ * \param mod A module containing definitions which can be referenced from
+ * \p expr. May be empty or undefined.
+ * \param expr An expression of function type to evaluate. May reference definitions from \p mod.
+ * \param device The device on which all primitives will be executed.
+ * \param target The compiler target flag for compiling primitives.
+ * \return A packed function that takes an array of Relay expressions and returns the
+ * result of applying \p expr to those arguments.
+ */
+TypedPackedFunc<ObjectRef(Array<Expr>)> EvalFunction(IRModule mod, Expr expr, Device device,
+                                                     Target target);
+
+/*!
+ * \brief Evaluates \p expr and returns its result.
+ *
+ * This function is intended to support TVM constant evaluation.
+ *
+ * \param expr An expression to evaluate.
+ * \param type_definitions Global type definitions which \p expr may references.
+ * \param import_set Already imported external modules.
+ * \param device The device on which all primitives will be executed.
+ * \param target The compiler target flag for compiling primitives.
+ * @return The object representing the result.
+ */
+ObjectRef Eval(Expr expr, Map<GlobalTypeVar, TypeData> type_definitions,
+               std::unordered_set<String> import_set, Device device, Target target);
+
 }  // namespace relay
 }  // namespace tvm
+
 #endif  // TVM_RELAY_INTERPRETER_H_

python/tvm/relay/analysis/analysis.py

@@ -433,8 +433,7 @@ def get_calibration_data(mod, data):
     mod = _ffi_api.get_calibrate_module(mod)
     mod = transform.Inline()(mod)
 
-    ref_ex = build_module.create_executor("graph", mod=mod, device=cpu(0))
-    ref_res = ref_ex.evaluate()(**data)
+    ref_res = build_module.create_executor("graph", mod=mod, device=cpu(0)).evaluate()(**data)
 
     calib_data = {}
     for gvar, indices in output_map.items():

python/tvm/relay/backend/interpreter.py

@@ -22,10 +22,9 @@
 
 import tvm._ffi
 from tvm.runtime import container, Object
-from tvm.ir import IRModule
 
 from . import _backend
-from .. import _make, analysis, transform
+from .. import _make, analysis
 from ... import nd
 from ..expr import Tuple, RefCreate, Call, Constant, GlobalVar, const
 from ..function import Function
@@ -178,6 +177,7 @@ def evaluate(self, expr=None, binds=None):
             return self._make_executor(expr)
 
         # normal expression evaluated by running a function.
+        # TODO(mbs): This should really be type rather than syntax driven.
         func = Function([], expr)
         return self._make_executor(func)()
 
@@ -196,65 +196,54 @@ class Interpreter(Executor):
 
     target : tvm.Target
         The target option to build the function.
+
+    CAUTION: Despite the API the module is prepared upon each call to evaluate
+    rather than once in create_executor.
+    That is:
+    .. code-block:: python
+
+        executor = relay.create_executor(kind="debug", mod=module)
+        a = executor.evaluate(expr)(args1)
+        b = executor.evaluate(expr)(args2)
+
+    will prepare all the bindings in module twice. For efficiency, try to hoist
+    calls to evaluate as high as possible, preferably immediately after create_executor:
+    .. code-block:: python
+
+        func = relay.create_executor(kind="debug", mod=module).evaluate(expr)
+        a = func(args1)
+        b = func(args2)
     """
 
     def __init__(self, mod, device, target):
         self.mod = mod
         self.device = device
         self.target = target
 
-    def optimize(self):
-        """Optimize functions in a module.
-
-        Returns
-        -------
-        opt_mod : tvm.IRModule
-            The optimized module.
-        """
-        seq = tvm.transform.Sequential(
-            [
-                # tvm.parser.AnnotateSpans(),
-                transform.SimplifyInference(),
-                transform.FuseOps(0),
-                transform.ToANormalForm(),
-                transform.InferType(),
-            ]
-        )
-        mod = seq(self.mod)
-        return mod
-
     def _make_executor(self, expr=None):
         if expr is None or isinstance(expr, GlobalVar):
             assert self.mod is not None
 
-        _intrp = _backend.CreateInterpreter(self.optimize(), self.device, self.target)
+        if expr is None:
+            # A missing expr denotes 'main' in the given module.
+            expr = self.mod.get_global_var("main")
 
-        def _interp_wrapper(*args, **kwargs):
-            if expr is None:
-                args = self._convert_args(self.mod["main"], args, kwargs)
+        # Evaluate expr to a packed function we can efficiently re-apply
+        # to Relay arguments.
+        func = _backend.EvalFunction(self.mod, expr, self.device, self.target)
+
+        def _apply_args(*args, **kwargs):
+            if isinstance(expr, GlobalVar):
+                # When expanding args, look inside the actual global definition so kwargs
+                # can be matched.
+                args = self._convert_args(self.mod[expr.name_hint], args, kwargs)
             else:
                 args = self._convert_args(expr, args, kwargs)
-
+            # Reflect python arguments up into Relay.
             relay_args = []
             for arg in args:
                 relay_args.append(_arg_to_ast(self.mod, arg))
+            # Apply func to Relay args
+            return func(relay_args)
 
-            # Set the entry function for the module.
-            if expr is None:
-                pass
-            elif isinstance(expr, GlobalVar):
-                self.mod["main"] = self.mod[expr]
-            else:
-                assert isinstance(expr, Function)
-                func = Function([], Call(expr, relay_args))
-                relay_args = []
-                if self.mod:
-                    self.mod["main"] = func
-                else:
-                    self.mod = IRModule.from_expr(func)
-
-            mod = self.optimize()
-            opt_expr = Call(mod["main"], relay_args)
-            return _intrp(opt_expr)
-
-        return _interp_wrapper
+        return _apply_args

python/tvm/relay/build_module.py

@@ -511,6 +511,9 @@ def _graph_wrapper(*args, **kwargs):
         return _graph_wrapper
 
 
+# TODO(mbs): Collapse the create_executor/evaluate phases together since a) most callers don't
+# reuse the executor for multiple expressions and b) any preparation necessary for the expression
+# evaluation needs to (currently) be done along with preparation for the module.
 def create_executor(kind="debug", mod=None, device=None, target="llvm", params=None):
     """Factory function to create an executor.
 

python/tvm/relay/frontend/common.py

@@ -545,11 +545,12 @@ def infer_value(input_val, params, mod=None):
             mod["main"] = _function.Function(analysis.free_vars(input_val), input_val)
         else:
             mod = IRModule.from_expr(input_val)
-        exc = tvm.relay.create_executor("debug", mod=mod, device=tvm.cpu(), target="llvm")
         inputs = []
         for param in mod["main"].params:
             inputs.append(params[param.name_hint])
-        result = exc.evaluate()(*inputs)
+        result = tvm.relay.create_executor(
+            "debug", mod=mod, device=tvm.cpu(), target="llvm"
+        ).evaluate()(*inputs)
         return result
 
 

python/tvm/relay/testing/__init__.py

@@ -134,10 +134,13 @@ def check_grad(
         test_inputs = inputs
 
     for target, dev in enabled_targets():
-        intrp = relay.create_executor(device=dev, target=target)
+        # Eval the backward and forward functions
+        # TODO(mbs): Evaluate a pair of functions so can share preparation between them.
+        bwd_func_compiled = relay.create_executor(device=dev, target=target).evaluate(bwd_func)
+        fwd_func_compiled = relay.create_executor(device=dev, target=target).evaluate(fwd_func)
 
         # Get analytic gradients.
-        _, grads = intrp.evaluate(bwd_func)(*inputs)
+        _, grads = bwd_func_compiled(*inputs)
         grads = [grad.numpy().astype("float64") for grad in grads]
 
         # Throw out gradients we aren't testing
@@ -154,7 +157,6 @@ def check_grad(
         assert len(grads) > 0, "You must test at least one gradient."
 
         # Get numeric gradients for each dimension of each param, using two-sided approximation.
-        fwd_func_compiled = intrp.evaluate(fwd_func)
         approx_grads = []
         for x in test_inputs:
             approx_grad = np.zeros(x.shape)