[Collage] CombinerRule and CandidatePartition::EstimateCost

src/relay/collage/candidate_function_cache.cc

@@ -0,0 +1,49 @@
+/*
+ * 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 src/relay/collage/candidate_function_cache.cc
+ * \brief A cache of the unique global name and costs for partitioned functions.
+ */
+
+#include "./candidate_function_cache.h"
+
+namespace tvm {
+namespace relay {
+namespace collage {
+
+CandidateFunctionCache::Entry& CandidateFunctionCache::GetEntry(const std::string& label,
+                                                                const Function& function) {
+  auto itr = cache_.find(function);
+  if (itr == cache_.end()) {
+    String compiler = function->GetAttr<String>(attr::kCompiler, String("tvm")).value();
+    std::string global_symbol_name = name_supply_->Fresh({compiler, label});
+    GlobalVar global_symbol(std::move(global_symbol_name), function->checked_type());
+    itr = cache_.emplace(function, Entry(std::move(global_symbol))).first;
+  }
+  return itr->second;
+}
+
+GlobalVar CandidateFunctionCache::GetGlobalSymbol(const Function& function) {
+  return GetEntry(/*label=*/"", function).global_symbol;
+}
+
+}  // namespace collage
+}  // namespace relay
+}  // namespace tvm

src/relay/collage/candidate_function_cache.h

@@ -0,0 +1,79 @@
+/*
+ * 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 src/relay/collage/candidate_function_cache.h
+ * \brief A cache of the unique global symbol name and cost for partitioned functions.
+ */
+
+#ifndef TVM_RELAY_COLLAGE_CANDIDATE_FUNCTION_CACHE_H_
+#define TVM_RELAY_COLLAGE_CANDIDATE_FUNCTION_CACHE_H_
+
+#include <tvm/relay/function.h>
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+#include "../transforms/compiler_function_utils.h"
+#include "./cost.h"
+#include "./name_supply.h"
+
+namespace tvm {
+namespace relay {
+namespace collage {
+
+/*!
+ * \brief A cache of the unique global symbol and cost for functions extracted to represent
+ * partitions. If two functions are structurally equal (which includes equality of their "Compiler"
+ * attributes) then they will share the same global symbol and estimated cost. We rely on the
+ * function's attributes to distinguish partitions which are structurally the same graph but
+ * intended for different targets.
+ */
+class CandidateFunctionCache : public transform::GlobalSymbolCache {
+ public:
+  explicit CandidateFunctionCache(std::shared_ptr<NameSupply> name_supply)
+      : name_supply_(std::move(name_supply)) {}
+
+  struct Entry {
+    GlobalVar global_symbol;
+    Cost cost = Cost::Unknown();  // Filled in when have estimated cost.
+
+    explicit Entry(GlobalVar global_symbol) : global_symbol(std::move(global_symbol)) {}
+  };
+
+  /*!
+   * \brief Returns the unique entry for \p function. If no such entry already exists, create it
+   * and assign it a unique global symbol name.
+   */
+  Entry& GetEntry(const std::string& label, const Function& function);
+
+  GlobalVar GetGlobalSymbol(const Function& function) final;
+
+ private:
+  std::shared_ptr<NameSupply> name_supply_;
+  std::unordered_map<Function, Entry, StructuralHash, StructuralEqual> cache_;
+};
+
+}  // namespace collage
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_COLLAGE_CANDIDATE_FUNCTION_CACHE_H_

src/relay/collage/candidate_partition.cc

@@ -24,8 +24,12 @@
 
 #include "./candidate_partition.h"
 
+#include <tvm/relay/analysis.h>
 #include <tvm/relay/attrs/memory.h>
+#include <tvm/relay/transform.h>
 
+#include "../transforms/compiler_function_utils.h"
+#include "./candidate_function_cache.h"
 #include "./candidate_set.h"
 #include "./partition_rule.h"
 #include "./partition_spec.h"
@@ -106,6 +110,102 @@ std::string CandidatePartitionNode::ToString() const {
   return os.str();
 }
 
+namespace {
+/*!
+ * \brief If function's body is a call to an inlined "Primitive" function, return it.
+ * Otherwise return function directly.
+ */
+Function GetPrimitiveFunction(const Function& function) {
+  if (const auto* call_node = function->body.as<CallNode>()) {
+    if (const auto* function_node = call_node->op.as<FunctionNode>()) {
+      if (function_node->HasNonzeroAttr(attr::kPrimitive)) {
+        return GetRef<Function>(function_node);
+      }
+    }
+  }
+  return function;
+}
+
+/*!
+ * \brief Eta-expand any tuple arguments of \p function. Ie rewrite:
+ * \code
+ *   f(x: (t1, t2)) { ... x ... }
+ * \endcode
+ * to
+ * \code
+ *   f(x_1: t1, x_2: t2) { ... (x_1, x_2) ... }
+ * \endcode
+ */
+Function EtaExpandTuples(const Function& function) {
+  Map<Var, Expr> subst;
+  Array<Var> new_params;
+  for (const auto& param : function->params) {
+    std::vector<TensorType> tensor_types = FlattenTupleType(param->type_annotation);
+    if (tensor_types.size() == 1) {
+      new_params.push_back(param);
+    } else {
+      Array<Expr> fields;
+      for (size_t i = 0; i < tensor_types.size(); ++i) {
+        Var new_param(param->name_hint() + "_" + std::to_string(i), tensor_types[i], param->span);
+        new_param->checked_type_ = tensor_types[i];
+        new_params.push_back(new_param);
+        fields.push_back(new_param);
+      }
+      Tuple new_tuple(fields);
+      subst.Set(param, new_tuple);
+    }
+  }
+  if (subst.empty()) {
+    return function;
+  }
+  return WithFields(function, new_params, Bind(function->body, subst));
+}
+
+}  // namespace
+
+Cost CandidatePartitionNode::EstimatedCost(
+    const DataflowGraph& dataflow_graph, const CostEstimator& cost_estimator,
+    const std::shared_ptr<CandidateFunctionCache>& cache) const {
+  if (cost_.is_unknown()) {
+    VLOG_CONTEXT << "spec " << partition_spec_name();
+    Function extracted_function = sub_graph_->ExtractAsFunction(dataflow_graph);
+    VLOG(2) << "Extracted function:" << std::endl << PrettyPrint(extracted_function);
+    extracted_function = EtaExpandTuples(extracted_function);
+    VLOG(2) << "Validating function:" << std::endl << PrettyPrint(extracted_function);
+    String error = partition_spec()->validate_sub_graph_func_(extracted_function);
+    if (!error.empty()) {
+      cost_ = Cost::Invalid();
+      VLOG(1) << "Unable to rewrite function: " << error;
+    } else {
+      // The extracted function may be the eta-expansion of a "Primitive" function.
+      // If so we want the cached external name and cost to be w.r.t. that function
+      // rather than the outer so that we'll get a cache hit when we outline functions
+      // in the final program.
+      Function primitive_function = GetPrimitiveFunction(extracted_function);
+      CandidateFunctionCache::Entry& entry =
+          cache->GetEntry(sub_graph_->label_, primitive_function);
+      if (entry.cost.is_unknown()) {
+        IRModule mod = IRModule::FromExpr(extracted_function);
+        VLOG(1) << "Outlining:" << std::endl << PrettyPrint(mod);
+        mod = OutlineCompilerFunctions(cache)(mod);
+        VLOG(1) << "Estimating cost of:" << std::endl
+                << PrettyPrint(mod) << std::endl
+                << "using target " << target()->ToDebugString();
+        entry.cost = cost_estimator->Estimate(mod, target(),
+                                              /*needs_tvm_tuning=*/!target().IsExternalCodegen());
+        VLOG(1) << "Measured cost as " << entry.cost.ToString();
+      } else {
+        VLOG(1) << "Reusing cost " << entry.cost.ToString()
+                << " cached in candidate function cache";
+      }
+      cost_ = entry.cost;
+    }
+  } else {
+    VLOG(1) << "Reusing cost " << cost_.ToString() << " cached in candidate";
+  }
+  return cost_;
+}
+
 CandidatePartition::CandidatePartition(String rule_name, SubGraph sub_graph,
                                        ObjectRef /* actually PartitionSpec */ spec, Cost cost) {
   auto node = runtime::make_object<CandidatePartitionNode>();

src/relay/collage/candidate_partition.h

@@ -32,7 +32,10 @@
 #include <string>
 #include <vector>
 
+#include "./candidate_function_cache.h"
 #include "./cost.h"
+#include "./cost_estimator.h"
+#include "./name_supply.h"
 #include "./sub_graph.h"
 
 namespace tvm {
@@ -93,6 +96,13 @@ class CandidatePartitionNode : public Object {
    */
   Target target() const;
 
+  /*!
+   * \brief Return the estimated cost of the candidate partition, using \p cost_estimator and
+   * \p cache.
+   */
+  Cost EstimatedCost(const DataflowGraph& dataflow_graph, const CostEstimator& cost_estimator,
+                     const std::shared_ptr<CandidateFunctionCache>& cache) const;
+
   /*!
    * \brief Returns a brief description of candidate suitable for debugging output.
    */