1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file src/relay/collage/cost_estimator.cc
22 * \brief Interface for measuring candidate partition cost.
23 */
24
25#ifndef TVM_RELAY_COLLAGE_COST_ESTIMATOR_H_
26#define TVM_RELAY_COLLAGE_COST_ESTIMATOR_H_
27
28#include <tvm/relay/function.h>
29
30#include "./cost.h"
31
32namespace tvm {
33namespace relay {
34namespace collage {
35
36/*!
37 * \brief An (abstract) estimator for the cost of executing "main" in an \p IRModule representing
38 * a candidate partition, using the given target for lowering and codegen.
39 *
40 * Generally the implementation will compile to a \p runtime::Module (possibly on a target-specific
41 * worker if cross-compilation is not available), repeatedly invoke "main" with random data until
42 * measure variance is acceptable (on a target-specific worker), and return the summarized costs.
43 *
44 * If using a TVM native \p Target, it is possible compilation will itself invoke TVM tuning.
45 *
46 * TODO(mbs): Actually, currently not abstract so can get some local measurements.
47 */
48class CostEstimatorNode : public Object {
49 public:
50 /*!
51 * \brief Returns the estimated cost (possibly after many many minutes of training time) of
52 * running "main" in \p mod using \p target, which represents a possible partitioning of
53 * some overall Relay expression.
54 */
55 virtual Cost Estimate(const IRModule& mod, const Target& target) const;
56
57 static constexpr const char* _type_key = "relay.collage.CostEstimator";
58 TVM_DECLARE_BASE_OBJECT_INFO(CostEstimatorNode, Object);
59};
60
61class CostEstimator : public ObjectRef {
62 public:
63 CostEstimator();
64 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CostEstimator, ObjectRef, CostEstimatorNode);
65};
66
67} // namespace collage
68} // namespace relay
69} // namespace tvm
70
71#endif // TVM_RELAY_COLLAGE_COST_ESTIMATOR_H_
72