tvm
builder.h
Go to the documentation of this file.
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_BUILDER_H_
20 #define TVM_META_SCHEDULE_BUILDER_H_
21 
22 #include <tvm/ffi/container/array.h>
23 #include <tvm/ffi/container/map.h>
24 #include <tvm/ffi/function.h>
25 #include <tvm/ffi/optional.h>
26 #include <tvm/ffi/reflection/registry.h>
27 #include <tvm/ffi/string.h>
28 #include <tvm/ir/module.h>
29 #include <tvm/runtime/ndarray.h>
30 #include <tvm/runtime/object.h>
31 #include <tvm/target/target.h>
32 
33 namespace tvm {
34 namespace meta_schedule {
35 
37 class BuilderInputNode : public runtime::Object {
38  public:
44  Optional<Map<String, runtime::NDArray>> params;
45 
46  static void RegisterReflection() {
47  namespace refl = tvm::ffi::reflection;
48  refl::ObjectDef<BuilderInputNode>()
49  .def_ro("mod", &BuilderInputNode::mod)
50  .def_ro("target", &BuilderInputNode::target)
51  .def_ro("params", &BuilderInputNode::params);
52  }
53 
54  static constexpr const char* _type_key = "meta_schedule.BuilderInput";
56 };
57 
62 class BuilderInput : public runtime::ObjectRef {
63  public:
70  TVM_DLL explicit BuilderInput(IRModule mod, Target target,
71  Optional<Map<String, runtime::NDArray>> params = std::nullopt);
73 };
74 
76 class BuilderResultNode : public runtime::Object {
77  public:
79  Optional<String> artifact_path;
81  Optional<String> error_msg;
82 
83  static void RegisterReflection() {
84  namespace refl = tvm::ffi::reflection;
85  refl::ObjectDef<BuilderResultNode>()
86  .def_ro("artifact_path", &BuilderResultNode::artifact_path)
87  .def_ro("error_msg", &BuilderResultNode::error_msg);
88  }
89 
90  static constexpr const char* _type_key = "meta_schedule.BuilderResult";
92 };
93 
98 class BuilderResult : public runtime::ObjectRef {
99  public:
105  TVM_DLL explicit BuilderResult(Optional<String> artifact_path, Optional<String> error_msg);
107 };
108 
110 class BuilderNode : public runtime::Object {
111  public:
113  virtual ~BuilderNode() = default;
119  virtual Array<BuilderResult> Build(const Array<BuilderInput>& build_inputs) = 0;
125  using FBuild = ffi::TypedFunction<Array<BuilderResult>(const Array<BuilderInput>&)>;
126 
127  static constexpr const char* _type_key = "meta_schedule.Builder";
129 };
130 
135 class Builder : public runtime::ObjectRef {
136  public:
144 };
145 
147 class PyBuilderNode : public BuilderNode {
148  public:
151 
152  static void RegisterReflection() {
153  namespace refl = tvm::ffi::reflection;
154  refl::ObjectDef<PyBuilderNode>().def_ro("f_build", &PyBuilderNode::f_build);
155  }
156 
157  Array<BuilderResult> Build(const Array<BuilderInput>& build_inputs) final {
158  ICHECK(f_build != nullptr) << "PyBuilder's Build method not implemented!";
159  return f_build(build_inputs);
160  }
161 
162  static constexpr const char* _type_key = "meta_schedule.PyBuilder";
164 };
165 
166 } // namespace meta_schedule
167 } // namespace tvm
168 
169 #endif // TVM_META_SCHEDULE_BUILDER_H_
Managed reference class to IRModuleNode.
Definition: module.h:257
Managed reference class to TargetNode.
Definition: target.h:191
The builder's input, containing an IRModule and the target.
Definition: builder.h:37
static constexpr const char * _type_key
Definition: builder.h:54
Optional< Map< String, runtime::NDArray > > params
Parameters for Relax build module.
Definition: builder.h:44
static void RegisterReflection()
Definition: builder.h:46
IRModule mod
The IRModule to be built.
Definition: builder.h:40
TVM_DECLARE_FINAL_OBJECT_INFO(BuilderInputNode, runtime::Object)
Target target
The target to be built for.
Definition: builder.h:42
Managed reference to BuilderInputNode.
Definition: builder.h:62
BuilderInput(IRModule mod, Target target, Optional< Map< String, runtime::NDArray >> params=std::nullopt)
Constructor of BuilderInput.
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode)
The abstract builder interface.
Definition: builder.h:110
virtual Array< BuilderResult > Build(const Array< BuilderInput > &build_inputs)=0
Generate the build results from build inputs.
ffi::TypedFunction< Array< BuilderResult >(const Array< BuilderInput > &)> FBuild
The function type of Build method.
Definition: builder.h:125
TVM_DECLARE_BASE_OBJECT_INFO(BuilderNode, runtime::Object)
static constexpr const char * _type_key
Definition: builder.h:127
virtual ~BuilderNode()=default
Default destructor.
The builder's output, containing the artifact path or error message if any.
Definition: builder.h:76
static void RegisterReflection()
Definition: builder.h:83
static constexpr const char * _type_key
Definition: builder.h:90
Optional< String > error_msg
The error message if any.
Definition: builder.h:81
Optional< String > artifact_path
The path to the built artifact.
Definition: builder.h:79
TVM_DECLARE_FINAL_OBJECT_INFO(BuilderResultNode, runtime::Object)
Managed reference to BuilderResultNode.
Definition: builder.h:98
BuilderResult(Optional< String > artifact_path, Optional< String > error_msg)
Constructor of BuilderResult.
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderResult, runtime::ObjectRef, BuilderResultNode)
Managed reference to BuilderNode.
Definition: builder.h:135
static Builder PyBuilder(BuilderNode::FBuild f_build)
Create a builder with customized build method on the python-side.
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Builder, runtime::ObjectRef, BuilderNode)
An abstract builder with customized build method on the python-side.
Definition: builder.h:147
TVM_DECLARE_FINAL_OBJECT_INFO(PyBuilderNode, BuilderNode)
Array< BuilderResult > Build(const Array< BuilderInput > &build_inputs) final
Generate the build results from build inputs.
Definition: builder.h:157
FBuild f_build
The packed function to the Build function.
Definition: builder.h:150
static constexpr const char * _type_key
Definition: builder.h:162
static void RegisterReflection()
Definition: builder.h:152
IRModule that holds the functions and type definitions.
Definition: repr_printer.h:91
tvm::PrimExpr mod(const tvm::PrimExpr &a, const tvm::PrimExpr &b)
Definition: broadcast.h:306
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A device-independent managed NDArray abstraction.
A managed object in the TVM runtime.
Compilation target object.