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/object.h>
30 #include <tvm/runtime/tensor.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  ffi::Optional<ffi::Map<ffi::String, runtime::Tensor>> 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  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.BuilderInput", BuilderInputNode,
54  runtime::Object);
55 };
56 
61 class BuilderInput : public runtime::ObjectRef {
62  public:
69  TVM_DLL explicit BuilderInput(
70  IRModule mod, Target target,
71  ffi::Optional<ffi::Map<ffi::String, runtime::Tensor>> params = std::nullopt);
73 };
74 
76 class BuilderResultNode : public runtime::Object {
77  public:
79  ffi::Optional<ffi::String> artifact_path;
81  ffi::Optional<ffi::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  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.BuilderResult", BuilderResultNode,
90  runtime::Object);
91 };
92 
97 class BuilderResult : public runtime::ObjectRef {
98  public:
104  TVM_DLL explicit BuilderResult(ffi::Optional<ffi::String> artifact_path,
105  ffi::Optional<ffi::String> error_msg);
108 };
109 
111 class BuilderNode : public runtime::Object {
112  public:
114  virtual ~BuilderNode() = default;
120  virtual ffi::Array<BuilderResult> Build(const ffi::Array<BuilderInput>& build_inputs) = 0;
126  using FBuild = ffi::TypedFunction<ffi::Array<BuilderResult>(const ffi::Array<BuilderInput>&)>;
127 
128  static constexpr const bool _type_mutable = true;
129  TVM_FFI_DECLARE_OBJECT_INFO("meta_schedule.Builder", BuilderNode, runtime::Object);
130 };
131 
136 class Builder : public runtime::ObjectRef {
137  public:
142  explicit Builder(ObjectPtr<BuilderNode> data) : ObjectRef(data) {
143  TVM_FFI_ICHECK(data != nullptr);
144  }
152 };
153 
155 class PyBuilderNode : public BuilderNode {
156  public:
159 
160  static void RegisterReflection() {
161  namespace refl = tvm::ffi::reflection;
162  refl::ObjectDef<PyBuilderNode>().def_ro("f_build", &PyBuilderNode::f_build);
163  }
164 
165  ffi::Array<BuilderResult> Build(const ffi::Array<BuilderInput>& build_inputs) final {
166  ICHECK(f_build != nullptr) << "PyBuilder's Build method not implemented!";
167  return f_build(build_inputs);
168  }
170 };
171 
172 } // namespace meta_schedule
173 } // namespace tvm
174 
175 #endif // TVM_META_SCHEDULE_BUILDER_H_
Managed reference class to IRModuleNode.
Definition: module.h:256
Managed reference class to TargetNode.
Definition: target.h:192
The builder's input, containing an IRModule and the target.
Definition: builder.h:37
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.BuilderInput", BuilderInputNode, runtime::Object)
static void RegisterReflection()
Definition: builder.h:46
ffi::Optional< ffi::Map< ffi::String, runtime::Tensor > > params
Parameters for Relax build module.
Definition: builder.h:44
IRModule mod
The IRModule to be built.
Definition: builder.h:40
Target target
The target to be built for.
Definition: builder.h:42
Managed reference to BuilderInputNode.
Definition: builder.h:61
BuilderInput(IRModule mod, Target target, ffi::Optional< ffi::Map< ffi::String, runtime::Tensor >> params=std::nullopt)
Constructor of BuilderInput.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BuilderInput, runtime::ObjectRef, BuilderInputNode)
The abstract builder interface.
Definition: builder.h:111
virtual ffi::Array< BuilderResult > Build(const ffi::Array< BuilderInput > &build_inputs)=0
Generate the build results from build inputs.
ffi::TypedFunction< ffi::Array< BuilderResult >(const ffi::Array< BuilderInput > &)> FBuild
The function type of Build method.
Definition: builder.h:126
virtual ~BuilderNode()=default
Default destructor.
TVM_FFI_DECLARE_OBJECT_INFO("meta_schedule.Builder", BuilderNode, runtime::Object)
static constexpr const bool _type_mutable
Definition: builder.h:128
The builder's output, containing the artifact path or error message if any.
Definition: builder.h:76
static void RegisterReflection()
Definition: builder.h:83
ffi::Optional< ffi::String > error_msg
The error message if any.
Definition: builder.h:81
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.BuilderResult", BuilderResultNode, runtime::Object)
ffi::Optional< ffi::String > artifact_path
The path to the built artifact.
Definition: builder.h:79
Managed reference to BuilderResultNode.
Definition: builder.h:97
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BuilderResult, runtime::ObjectRef, BuilderResultNode)
BuilderResult(ffi::Optional< ffi::String > artifact_path, ffi::Optional< ffi::String > error_msg)
Constructor of BuilderResult.
Managed reference to BuilderNode.
Definition: builder.h:136
Builder(ObjectPtr< BuilderNode > data)
Constructor from ObjectPtr<BuilderNode>.
Definition: builder.h:142
static Builder PyBuilder(BuilderNode::FBuild f_build)
Create a builder with customized build method on the python-side.
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Builder, runtime::ObjectRef, BuilderNode)
An abstract builder with customized build method on the python-side.
Definition: builder.h:155
FBuild f_build
The packed function to the Build function.
Definition: builder.h:158
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("meta_schedule.PyBuilder", PyBuilderNode, BuilderNode)
ffi::Array< BuilderResult > Build(const ffi::Array< BuilderInput > &build_inputs) final
Generate the build results from build inputs.
Definition: builder.h:165
static void RegisterReflection()
Definition: builder.h:160
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:308
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
A managed object in the TVM runtime.
A device-independent managed Tensor abstraction.
Compilation target object.