[Unity][Analysis] Add utility for collecting compile-time bindings (#16312)

Whether an optimizations should be performed may depend on when the
variables in an expression are known.

For example, consider a LoRA-adjusted model, with base weights `W` of
shape `[m,n]`, LoRA components `A` and `B` with shapes `[r,n]` and
`[m,r]` respectively, and activations `x` with shape `[n,1]`.  The
LoRA-adjusted matmul could be computed either as `(W + B*A)*x` or as
`(W*x + B*(A*x))`.

If `A` and `B` are provided at run-time, then computing `(W +
B*(A*x))` requires significantly fewer computations.

* `(W + B*A)*x`: `m*n*(2*r + 3)` operations
  1. `B*A`: `2*m*n*r` operations using a naive matmul
  2. Adding `W` to (1): `m*n` operations
  3. Multiplying `x` by (2): `2*m*n` operations

* `(W*x + B*(A*x))`: (2*m*n + r*(2*n + 2*m + 1))
  1. `W*x`: `2*m*n` operations
  2. `A*x`: `2*r*n` operations
  3. Multiplying `B` by (2): `2*m*r` operations
  4. Adding (1) and (3)`: `m` operations

However, if `A` and `B` are known at compile-time, then computing `(W
+ B*A)*x` groups all compile-time values together, allowing them to be
computed earlier (i.e. using `LiftTransformParams`)

* `(W + B*A)*x`: `2*m*n` operations
  1. `B*A`: 0 operations, computed at compile-time
  2. Adding `W` to (1): 0 operations, computed at compile-time
  3. Multiplying `x` by (2): `2*m*n` operations

Since the choice of optimized expression depends on which parameters
can be computed at compile-time, it is useful to have a utility that
identifies values that can be computed at compile-time.

Author: Lunderberg

include/tvm/relax/analysis.h

@@ -533,6 +533,21 @@ TVM_DLL bool WellFormed(IRModule m, bool check_struct_info = true);
 TVM_DLL Map<tir::Block, Map<ObjectRef, tir::IndexMap>> SuggestLayoutTransforms(
     const Function& fn, Array<tir::IndexMap> write_buffer_transformations);
 
+/* \brief Collect variables whose value can be computed at compile-time
+ *
+ * If a function has the `kNumInput` attribute, then the first
+ * `kNumInput` parameters are provided at run-time, while all
+ * remaining parameters may be known at compile-time.  This utility
+ * collects all variable bindings that only depend, directly or
+ * indirectly, on the parameters known at compile-time.
+ *
+ * \param func The relax::Function to analyze
+ *
+ * \return The set of variables that can be computed at compile-time,
+ * in order of their occurrence within the function.
+ */
+TVM_DLL Array<Var> ComputableAtCompileTime(const Function& func);
+
 }  // namespace relax
 }  // namespace tvm
 

python/tvm/relax/analysis/__init__.py

@@ -21,6 +21,7 @@
     all_global_vars,
     all_vars,
     bound_vars,
+    computable_at_compile_time,
     contains_impure_call,
     definable_tir_vars_in_struct_info,
     defined_symbolic_vars,

python/tvm/relax/analysis/analysis.py

@@ -528,3 +528,28 @@ def detect_recursion(mod: tvm.IRModule) -> List[List[GlobalVar]]:
         with any other, it will be a singleton in this list.
     """
     return _ffi_api.detect_recursion(mod)  # type: ignore
+
+
+def computable_at_compile_time(func: Function) -> List[Var]:
+    """Collect variables whose value can be computed at compile-time
+
+    If a function has the `kNumInput` attribute, then the first
+    `kNumInput` parameters are provided at run-time, while all
+    remaining parameters may be known at compile-time.  This utility
+    collects all variable bindings that only depend, directly or
+    indirectly, on the parameters known at compile-time.
+
+    Parameters
+    ----------
+    func: Function
+
+        The `relax.Function` to analyze
+
+    Returns
+    -------
+    ret: List[Var]
+
+        The set of variables that can be computed at compile-time, in
+        order of their occurrence within the function.
+    """
+    return _ffi_api.computable_at_compile_time(func)  # type: ignore

src/relax/analysis/computable_at_compile_time.cc

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file computable_at_compile_time.cc
+ *
+ * \brief Utilities for identifying potentially compile-time variables
+ */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/expr_functor.h>
+
+#include "../../support/ordered_set.h"
+
+namespace tvm {
+namespace relax {
+
+namespace {
+class CompileTimeCollector : ExprVisitor {
+ public:
+  static Array<Var> Collect(const Function& func) {
+    CompileTimeCollector visitor;
+    visitor(func);
+    return Array<Var>(visitor.known_relax_vars_.begin(), visitor.known_relax_vars_.end());
+  }
+
+ private:
+  void VisitExpr_(const FunctionNode* func) override {
+    if (auto opt_num_input = func->attrs.GetAttr<Integer>(attr::kNumInput)) {
+      size_t num_input = opt_num_input.value()->value;
+      for (size_t i = num_input; i < func->params.size(); i++) {
+        MarkAsKnown(func->params[i]);
+      }
+    }
+
+    ExprVisitor::VisitExpr_(func);
+  }
+
+  void VisitBinding(const Binding& binding) override {
+    Expr value = GetBoundValue(binding);
+    bool can_compute_at_compile_time = [&]() {
+      for (const auto& relax_var : FreeVars(value)) {
+        if (!known_relax_vars_.count(relax_var)) {
+          return false;
+        }
+      }
+      for (const auto& tir_var : FreeSymbolicVars(value)) {
+        if (!known_tir_vars_.count(tir_var)) {
+          return false;
+        }
+      }
+
+      return true;
+    }();
+
+    if (can_compute_at_compile_time) {
+      MarkAsKnown(binding->var);
+    }
+
+    ExprVisitor::VisitBinding(binding);
+  }
+
+  void MarkAsKnown(const Var& var) {
+    known_relax_vars_.insert(var);
+    for (const auto& tir_var : DefinableTIRVarsInStructInfo(GetStructInfo(var))) {
+      known_tir_vars_.insert(tir_var);
+    }
+  }
+
+  support::OrderedSet<Var> known_relax_vars_;
+  std::unordered_set<tir::Var, ObjectPtrHash, ObjectPtrEqual> known_tir_vars_;
+};
+}  // namespace
+
+Array<Var> ComputableAtCompileTime(const Function& func) {
+  return CompileTimeCollector::Collect(func);
+}
+
+TVM_REGISTER_GLOBAL("relax.analysis.computable_at_compile_time")
+    .set_body_typed(ComputableAtCompileTime);
+
+}  // namespace relax
+}  // namespace tvm