[COLLAGE] Add more customization to support more targets (#13450)

* [COLLAGE] Add more customization to support more targets

1. Added custom cost module to provide a provision to incorporate
custom cost estimator python function instead using default cost function.
   eg: cost_estimator = CustomCostEstimator(py_fn_estimator="tvm.relay.collage.opencl_cost_estimator")
       mod = CollagePartition(config, cost_estimator=cost_estimator)(mod)
2. Added provision to select BYOC fusion style for all compiler target.
   eg : config = { "relay.collage.byoc_fusion_style": ["compiler.NoFusion", "compiler.TVMFusion"]}
        ctxt = tvm.transform.PassContext(config=config)

* Fix the lint errors

* Fix the lint error whitespace

* Fix the lint error tabs

* Fix the lint error tabs

* Fix the lint error tabs

* move the clml collage test case to test_clml

* Fix lint error whitespace

* Fix the import error

* Fix the envirnoment var and import

* Add comments

* Add clml preprocess module in cost estimator

* Fix whitespace lint error

* Fix whitespace lint error

* Fix whitespace lint error

* Fix the comments and removed unwanted code

* Fix whitespace error

* Removed Todo comments

* Removed TODO comments

* Updated naming convension

* Fix typo error

* Fixe the typo error

* Corrected typo error

* Corrected typo error

* Removed unused and fix typo error

* Removed redundent code and optimize the code

* Fix the lint error

* Fix whitespace lint error

* Removed Prints in file

* Fix lint error

* Fix lint error

* Removed runner template in test script

* Fix the lint error

* Fix lint error

* Fix lint error

* Fix the lint error

* Fix the lint error

Co-authored-by: kvegiraj <[email protected]>

Author: krishnaraj36

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_custom"):
+        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

@@ -23,6 +23,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 tvm.relay.expr_functor import ExprMutator
 from tvm.relay.expr import Call, TupleGetItem
 
@@ -161,6 +162,25 @@ def alter_conv(attrs, inputs, tinfos, out_type):
     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,60 @@
+/*
+ * 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 {
+  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,67 @@
+/*
+ * 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 target-specific cost functions.
+ */
+
+#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 cost estimator that uses a target-specific cost function.
+ */
+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 implemented cost function name.
+   */
+  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,9 @@ PartitionRule MakeTVMPartitionRule() {
 }
 
 /*!
- * \brief Returns the fusion style for \p compiler.
- *
- * TODO(mbs): Defer to per-BYOC integration definition.
+ * \brief Returns the fusion style for default compiler.
  */
-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 +101,35 @@ BYOCStyle BYOCFusionStyleForCompiler(const String& compiler) {
   }
 }
 
+/*!
+ * \brief Returns the fusion style for given compiler.
+ */
+BYOCStyle BYOCFusionStyleForCompiler(const String& compiler) {
+  tvm::transform::PassContext ctxt = tvm::transform::PassContext::Current();
+  std::string config_key = "relay.collage.byoc_fusion_style";
+  Optional<Array<String>> byoc_configs = ctxt->GetConfig(config_key, Optional<Array<String>>());
+  BYOCStyle byoc_fusion_style = DefaultBYOCFusionStyleForCompiler(compiler);
+  if (!byoc_configs) {
+    return byoc_fusion_style;
+  }
+  for (auto config_ : byoc_configs.value()) {
+    std::vector<std::string> byoc_cfg = SplitString(config_, ".");
+    if (byoc_cfg[0] == compiler) {
+      if (byoc_cfg[1] == "NoFusion") {
+        byoc_fusion_style = kNoFusionBYOCStyle;
+      } else if (byoc_cfg[1] == "TVMFusion") {
+        byoc_fusion_style = kTVMFusionBYOCStyle;
+      } else if (byoc_cfg[1] == "ArbitraryFusion") {
+        byoc_fusion_style = kArbitraryFusionBYOCStyle;
+      } else {
+        ICHECK(false) << "Invalid fusion name for compiler " << byoc_cfg[0] << " 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.

src/relay/collage/utils.cc

@@ -134,6 +134,19 @@ bool MustBeLowered(const Expr& expr) {
   return false;
 }
 
+std::vector<std::string> SplitString(std::string stmt, const char* del) {
+  std::vector<std::string> str_tokens;
+  int start = 0;
+  int end = stmt.find(del, 0);
+  str_tokens.emplace_back(stmt.substr(start, end));
+  while (end != -1) {
+    stmt = stmt.substr(end + 1, stmt.size());
+    end = stmt.find(del, 0);
+    str_tokens.emplace_back(stmt.substr(start, end));
+  }
+  return str_tokens;
+}
+
 }  // namespace collage
 }  // namespace relay
 }  // namespace tvm

src/relay/collage/utils.h

@@ -31,6 +31,7 @@
 #include <tvm/runtime/container/string.h>
 
 #include <string>
+#include <vector>
 
 namespace tvm {
 namespace relay {
@@ -79,6 +80,11 @@ bool IsSpecialOp(const OpNode* op_node);
  */
 bool MustBeLowered(const Expr& expr);
 
+/*!
+ * \brief Returns the list of split strings of given statement with delimiter.
+ */
+std::vector<std::string> SplitString(std::string stmt, const char* del);
+
 }  // namespace collage
 }  // namespace relay
 }  // namespace tvm