[COLLAGE] Add more customization to support more targets

python/tvm/relay/collage/__init__.py

@@ -21,4 +21,5 @@
     WARMUP_MIN_REPEAT_MS,
     CostEstimator,
     MockCostEstimator,
+    CustomCostEstimator,
 )

python/tvm/relay/collage/collage.py

@@ -52,6 +52,14 @@ def __init__(self, target_costs, max_estimates=0):
         self.__init_handle_by_constructor__(_ffi_api.MockCostEstimator, target_costs, max_estimates)
 
 
+@register_object("relay.collage.CustomCostEstimator")
+class CustomCostEstimator(Object):
+    """CustomEstimator class"""
+
+    def __init__(self, py_fn_estimator="tvm.relay.collage.estimate_seconds_cust"):
+        self.__init_handle_by_constructor__(_ffi_api.CustomCostEstimator, py_fn_estimator)
+
+
 def arg_for(arg_type, device):
     """Returns a test argument of Relay arg_type on device"""
     assert isinstance(arg_type, tvm.ir.TensorType)

python/tvm/relay/op/contrib/clml.py

@@ -22,7 +22,7 @@
 from tvm._ffi import register_func
 from tvm.relay import transform
 from tvm.relay.build_module import bind_params_by_name
-
+from tvm.relay import function as _function
 from ...dataflow_pattern import wildcard, is_op, is_constant, is_tuple_get_item, is_tuple
 from .register import register_pattern_table
 from ..strategy.generic import is_depthwise_conv2d
@@ -127,6 +127,25 @@ def convert_conv(attrs, inputs, tinfos, desired_layouts):
     return preprocessed_mod
 
 
+def preprocess_for_clml(mod):
+    """Preprocessing pass to alter the layouts for CLML compiler target"""
+
+    for _var in mod.get_global_vars():
+        if _var.name_hint == "main":
+            continue
+        fn = mod[_var.name_hint]
+        if "Compiler" in fn.attrs.keys() and fn.attrs["Compiler"] == "clml":
+            new_fn = fn.body
+            clml_mod = tvm.IRModule.from_expr(new_fn)
+            with tvm.transform.PassContext(opt_level=3):
+                clml_mod = preprocess_module(clml_mod)
+            new_body = clml_mod["main"].body
+            mod[_var.name_hint] = _function.Function(
+                fn.params, new_body, fn.ret_type, fn.type_params, fn.attrs
+            )
+    return mod
+
+
 @register_pattern_table("clml")
 def clml_pattern_table():
     """Get the CLML pattern table."""

src/relay/collage/collage_partitioner.cc

@@ -55,7 +55,7 @@ namespace {
 
 TVM_REGISTER_PASS_CONFIG_OPTION("relay.collage.tvm_max_depth", Integer);
 TVM_REGISTER_PASS_CONFIG_OPTION("relay.collage.byoc_max_depth", Integer);
-
+TVM_REGISTER_PASS_CONFIG_OPTION("relay.collage.byoc_fusion_style", Array<String>);
 /*!
  * \brief Represents the overall expression after some number of non-overlapping candidate
  * partitions have been applied.

src/relay/collage/custom_cost_estimator.cc

@@ -0,0 +1,61 @@
+/*
+ * 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/custom_cost_estimator.cc
+ * \brief A custom CostEstimator to support Alternative cost functions.
+ */
+
+#include "./custom_cost_estimator.h"
+
+#include <tvm/relay/expr_functor.h>
+
+namespace tvm {
+namespace relay {
+namespace collage {
+
+TVM_REGISTER_OBJECT_TYPE(CustomCostEstimatorNode);
+
+Cost CustomCostEstimatorNode::Estimate(const IRModule& mod, const Target& target) const {
+  // TODO(mbs): Eventually should be abstract. For now bounce to the Python local impl.
+  static const runtime::PackedFunc* estimate_seconds = runtime::Registry::Get(py_fn_estimator_);
+  ICHECK(estimate_seconds);
+  const double value = (*estimate_seconds)(mod, target);
+  if (std::isinf(value)) {
+    return Cost::Invalid();
+  } else if (std::isnan(value)) {
+    return Cost::Unknown();
+  } else {
+    return Cost::Value(value);
+  }
+}
+
+CustomCostEstimator::CustomCostEstimator(String py_fn_estimator) {
+  auto node = make_object<CustomCostEstimatorNode>();
+  node->py_fn_estimator_ = std::move(py_fn_estimator);
+  data_ = std::move(node);
+}
+
+TVM_REGISTER_GLOBAL("relay.collage.CustomCostEstimator").set_body_typed([](String py_fn_estimator) {
+  return CustomCostEstimator(std::move(py_fn_estimator));
+});
+
+}  // namespace collage
+}  // namespace relay
+}  // namespace tvm

src/relay/collage/custom_cost_estimator.h

@@ -0,0 +1,68 @@
+/*
+ * 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/custom_cost_estimator.cc
+ * \brief A custom CostEstimator to support all different compiler cost.
+ */
+
+#ifndef TVM_RELAY_COLLAGE_CUSTOM_COST_ESTIMATOR_H_
+#define TVM_RELAY_COLLAGE_CUSTOM_COST_ESTIMATOR_H_
+
+#include <tvm/relay/function.h>
+
+#include "./cost.h"
+#include "./cost_estimator.h"
+
+namespace tvm {
+namespace relay {
+namespace collage {
+
+/*!
+ * \brief A custom cost estimator which can determine the cost of a candidate based on
+ * the candidate's target.
+ */
+class CustomCostEstimatorNode : public CostEstimatorNode {
+ public:
+  Cost Estimate(const IRModule& mod, const Target& target) const override;
+
+  static constexpr const char* _type_key = "relay.collage.CustomCostEstimator";
+  TVM_DECLARE_FINAL_OBJECT_INFO(CustomCostEstimatorNode, CostEstimatorNode);
+
+ protected:
+  /*!
+   * \brief Python implementation function register name to invote.
+   */
+  String py_fn_estimator_;
+
+  friend class CustomCostEstimator;
+};
+
+class CustomCostEstimator : public CostEstimator {
+ public:
+  explicit CustomCostEstimator(String py_fn_estimator);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(CustomCostEstimator, CostEstimator, CustomCostEstimatorNode);
+};
+
+}  // namespace collage
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_COLLAGE_CUSTOM_COST_ESTIMATOR_H_

src/relay/collage/gather_partition_specs.cc

@@ -89,11 +89,11 @@ PartitionRule MakeTVMPartitionRule() {
 }
 
 /*!
- * \brief Returns the fusion style for \p compiler.
+ * \brief Returns the fusion style for default compiler.
  *
  * TODO(mbs): Defer to per-BYOC integration definition.
  */
-BYOCStyle BYOCFusionStyleForCompiler(const String& compiler) {
+BYOCStyle DefaultBYOCFusionStyleForCompiler(const String& compiler) {
   if (compiler == "cutlass" || compiler == "cublas" || compiler == "cudnn") {
     return kNoFusionBYOCStyle;
   } else if (compiler == "tensorrt") {
@@ -103,6 +103,40 @@ BYOCStyle BYOCFusionStyleForCompiler(const String& compiler) {
   }
 }
 
+/*!
+ * \brief Returns the fusion style for /p compiler.
+ *
+ * TODO(mbs): Defer to per-BYOC integration definition.
+ */
+BYOCStyle BYOCFusionStyleForCompiler(const String& compiler) {
+  tvm::transform::PassContext ctxt = tvm::transform::PassContext::Current();
+  std::string config_key = "relay.collage.byoc_fusion_style";
+  std::string compiler_id = compiler;
+  Optional<Array<String>> byoc_configs = ctxt->GetConfig(config_key, Optional<Array<String>>());
+  if (!byoc_configs.defined()) {
+    return DefaultBYOCFusionStyleForCompiler(compiler);
+  }
+  BYOCStyle byoc_fusion_style = kNoFusionBYOCStyle;
+  for (auto config_ : byoc_configs.value()) {
+    std::string byoc_str = static_cast<std::string>(config_);
+    std::string byoc_compiler = byoc_str.substr(0, byoc_str.find(".", 0));
+    if (byoc_compiler == compiler) {
+      std::string fusion_name = byoc_str.substr(byoc_str.find(".", 0) + 1, byoc_str.size());
+      if (fusion_name == "NoFusion") {
+        byoc_fusion_style = kNoFusionBYOCStyle;
+      } else if (fusion_name == "TVMFusion") {
+        byoc_fusion_style = kTVMFusionBYOCStyle;
+      } else if (fusion_name == "MaxDepthFusion") {
+        byoc_fusion_style = kArbitraryFusionBYOCStyle;
+      } else {
+        ICHECK(false) << "Invalid fusion name for compiler " << byoc_compiler << " in pass context";
+      }
+      break;
+    }
+  }
+  return byoc_fusion_style;
+}
+
 /*!
  * \brief Returns the primitive combiner rules which allow for any touching candidates
  * to be fused provided they don't have kind \p kOpaque.