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#ifndef TVM_META_SCHEDULE_EXTRACTED_TASK_H_
20#define TVM_META_SCHEDULE_EXTRACTED_TASK_H_
21
22#include <tvm/ir/module.h>
23#include <tvm/node/reflection.h>
24#include <tvm/runtime/container/array.h>
25#include <tvm/runtime/container/string.h>
26#include <tvm/runtime/object.h>
27#include <tvm/target/target.h>
28
29namespace tvm {
30namespace tir {
31class PrimFunc;
32} // namespace tir
33namespace te {
34class Tensor;
35} // namespace te
36} // namespace tvm
37
38namespace tvm {
39namespace meta_schedule {
40
41/*! \brief A tuning task extracted from the high-level IR */
42class ExtractedTaskNode : public runtime::Object {
43 public:
44 /*! \brief The name of the task extracted */
45 String task_name;
46 /*! \brief The high-level IR */
47 IRModule mod;
48 /*! \brief Target */
49 Target target;
50 /*! \brief A list of low-level IRs that the high-level IR could potentially dispatch to */
51 Array<IRModule> dispatched;
52 /*! \brief Weight of the task */
53 int weight;
54
55 void VisitAttrs(tvm::AttrVisitor* v) {
56 v->Visit("task_name", &task_name);
57 v->Visit("mod", &mod);
58 v->Visit("target", &target);
59 v->Visit("dispatched", &dispatched);
60 v->Visit("weight", &weight);
61 }
62
63 static constexpr const char* _type_key = "meta_schedule.ExtractedTask";
64 TVM_DECLARE_FINAL_OBJECT_INFO(ExtractedTaskNode, runtime::Object);
65};
66
67/*!
68 * \brief Managed reference to ExtractedTaskNode
69 * \sa ExtractedTaskNode
70 */
71class ExtractedTask : public runtime::ObjectRef {
72 public:
73 explicit ExtractedTask(String task_name, IRModule mod, Target target, Array<IRModule> dispatched,
74 int weight);
75 TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(ExtractedTask, runtime::ObjectRef,
76 ExtractedTaskNode);
77};
78
79} // namespace meta_schedule
80} // namespace tvm
81
82#endif // TVM_META_SCHEDULE_EXTRACTED_TASK_H_
83