- Andrew's comments

Author: mbs-octoml

include/tvm/relay/function.h

@@ -172,23 +172,29 @@ const FunctionNode* AsOptimizableFunctionNode(const BaseFunc& base_func);
 namespace attr {
 
 /*!
- * \brief Mark the function as a primitive function. Should be bound to \p Integer(1).
+ * \brief Mark the function as representing a sub-graph which is to be lowered or compiled as
+ * a unit. For example, the function may represent a kernel which TVM will lower to a PrimFunc.
+ * If present should be bound to \p Integer(1). May be accompanied by "Compiler", see below.
+ * The function body should be considered opaque by Relay, and many passes simply ignore these
+ * functions.
  *
  * Type: Integer
  */
 constexpr const char* kPrimitive = "Primitive";
 
 /*!
- * \brief Mark the function as being 'extern', ie implemented in a runtime::Module. Should be bound
- * to \p Integer(1). Typically accompanied by "Primitive".
+ * \brief Mark the function as externally implemented, ie bound in a runtime::Module within the
+ * IRModule's "external_mods" attribute. If present should be bound to \p Integer(1). Generally
+ * the only attribute when present.
  *
  * Type: Integer
  */
 constexpr const char* kExtern = "Extern";
 
 /*!
- * \brief Indicate the external codegen 'compiler' that should be used for building this function.
- * When this is unset or set to "default", the default compilation pipeline will be used.
+ * \brief Indicates the name of the external codegen 'compiler' that should be used to lower
+ * or compile the function other than TVM's default lowering pipeline. The name may correspond
+ * to a TargetKind name. There may be a global function registered under 'relay.ext.{name}'.
  *
  * Type: String
  */

python/tvm/relay/transform/transform.py

@@ -1371,8 +1371,10 @@ def SplitArgs(max_function_args):
 
 
 def OutlineCompilerFunctionsWithExistingGlobalSymbols(compiler_filter=""):
-    """A pass to outline all literal functions in direct call positions which have a "Compiler"
-    attribute. The functions are bound to unique global vars according to their existing
+    """Outlines all literal functions in direct call positions which have a "Compiler"
+    attribute.
+
+    The outlined functions are bound to unique global vars according to their existing
     "global_symbol" attribute. At most one function with the same global symbol is outlined.
 
     If compiler_filter is non-empty only functions with that as their attribute value are
@@ -1395,9 +1397,11 @@ def OutlineCompilerFunctionsWithExistingGlobalSymbols(compiler_filter=""):
 
 
 def MarkCompilerFunctionsAsExtern(compiler_filter=""):
-    """A pass to mark all global functions which have a "Compiler" attribute matching
-    compiler_filter as 'extern' by replacing all attributes with a single "Extern" attribute, and
-    rewrite all calls to such functions to use the 'call_lowered' calling convention.
+    """Marks all global functions which have a "Compiler" attribute matching
+    compiler_filter as 'extern'.
+
+    The function's attributes are replaced with a single "Extern" attribute, and
+    all calls to the function are switched to use the 'call_lowered' calling convention.
 
     If compiler_filter is non-empty only functions with that as their attribute value are
     outlined.

src/relay/backend/te_compiler.cc

@@ -176,7 +176,7 @@ class TECompilerImpl : public TECompilerNode {
                 WithFields(GetRef<Function>(function_node), function_node->params,
                            function_node->body, function_node->ret_type, function_node->type_params,
                            /* erase attributes */ DictAttrs(Map<String, ObjectRef>()));
-            // Mark function as 'extern' using the "ExternalSymbol" attribute.
+            // Mark function as 'extern'.
             function = WithAttr(std::move(function), attr::kExtern, Integer(1));
             module->Add(kv2.first, function);
           }

src/relay/transforms/compiler_function_utils.cc

@@ -49,8 +49,9 @@ class Outliner : public MixedModeMutator {
       if (opt_compiler.defined() &&
           (compiler_filter_.empty() || opt_compiler.value() == compiler_filter_)) {
         auto function = GetRef<Function>(function_node);
-        ICHECK(FreeVars(function).empty()) << "Function marked with '" << attr::kCompiler
+        DCHECK(FreeVars(function).empty()) << "Function marked with '" << attr::kCompiler
                                            << "' attribute should not have free variables";
+        // Ask the cache to supply a unique  global var for this function.
         GlobalVar global_symbol = cache_->GetGlobalSymbol(function);
         // Depending on the cache's implementation, two structurally equal (but not object equal)
         // functions may be assigned the same global symbol. If so we'll lift it just once, but
@@ -60,15 +61,23 @@ class Outliner : public MixedModeMutator {
               WithAttr(std::move(function), tvm::attr::kGlobalSymbol, global_symbol->name_hint);
           mod_->Add(global_symbol, function);
         }
+        // Update the call.
         return WithFields(new_call, global_symbol);
       }
     }
     return post;
   }
 
  private:
+  /*!
+   * \brief A cached mapping from functions to global variables. Depending on the implementation
+   * the cache may generate fresh symbols or require the function to already have a "global_symbol"
+   * attribute, and may share symbols between structurally equal functions.
+   */
   GlobalSymbolCache* cache_;
+  /*! \brief If non-empty, the "Compiler" attribute value to require on functions to outline. */
   std::string compiler_filter_;
+  /*! \brief Module being rewritten. */
   IRModule mod_;
 };
 
@@ -102,7 +111,9 @@ class CallRewriter : public MixedModeMutator {
   }
 
  private:
+  /*! \brief If non-empty, the "Compiler" attribute value to require on functions to outline. */
   std::string compiler_filter_;
+  /*! \brief Module being rewritten. */
   IRModule mod_;
 };