Skip to content

Commit 5ff8a9e

Browse files
junrushaozxybazhspectrometerHBHMasterJH5574jinhongyii
authored and
ylc
committed
[Meta Schedule][M3b] Builder (apache#9044)
* [Meta Schedule][M3b] Builder This PR is part of the meta schedule project (apache#8473) Co-authored-by: Xiyou Zhou <[email protected]> Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]> * add typing * unreachable Co-authored-by: Xiyou Zhou <[email protected]> Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]>
1 parent e67536a commit 5ff8a9e

File tree

13 files changed

+992
-3
lines changed

13 files changed

+992
-3
lines changed

‎CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,7 @@ assign_source_group("Include" ${GROUP_INCLUDE})
257257
# Source file lists
258258
file(GLOB_RECURSE COMPILER_SRCS
259259
src/auto_scheduler/*.cc
260+
src/meta_schedule/*.cc
260261
src/node/*.cc
261262
src/ir/*.cc
262263
src/arith/*.cc

‎include/tvm/meta_schedule/builder.h‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
2+
/*
3+
* Licensed to the Apache Software Foundation (ASF) under one
4+
* or more contributor license agreements. See the NOTICE file
5+
* distributed with this work for additional information
6+
* regarding copyright ownership. The ASF licenses this file
7+
* to you under the Apache License, Version 2.0 (the
8+
* "License"); you may not use this file except in compliance
9+
* with the License. You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing,
14+
* software distributed under the License is distributed on an
15+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16+
* KIND, either express or implied. See the License for the
17+
* specific language governing permissions and limitations
18+
* under the License.
19+
*/
20+
#ifndef TVM_META_SCHEDULE_BUILDER_H_
21+
#define TVM_META_SCHEDULE_BUILDER_H_
22+
23+
#include <tvm/ir/module.h>
24+
#include <tvm/target/target.h>
25+
26+
namespace tvm {
27+
namespace meta_schedule {
28+
29+
/*! \brief The builder's input. */
30+
class BuilderInputNode : public runtime::Object {
31+
public:
32+
/*! \brief The IRModule to be built. */
33+
IRModule mod;
34+
/*! \brief The target to be built for. */
35+
Target target;
36+
37+
void VisitAttrs(tvm::AttrVisitor* v) {
38+
v->Visit("mod", &mod);
39+
v->Visit("target", &target);
40+
}
41+
42+
static constexpr const char* _type_key = "meta_schedule.BuilderInput";
43+
TVM_DECLARE_FINAL_OBJECT_INFO(BuilderInputNode, runtime::Object);
44+
};
45+
46+
/*!
47+
* \brief Managed reference to BuilderInputNode
48+
* \sa BuilderInputNode
49+
*/
50+
class BuilderInput : public runtime::ObjectRef {
51+
public:
52+
/*!
53+
* \brief Constructor of BuilderInput.
54+
* \param mod The IRModule to be built.
55+
* \param target The target to be built for.
56+
*/
57+
TVM_DLL explicit BuilderInput(IRModule mod, Target target);
58+
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode);
59+
};
60+
61+
/*! \brief The builder's output. */
62+
class BuilderResultNode : public runtime::Object {
63+
public:
64+
/*! \brief The path to the built artifact. */
65+
Optional<String> artifact_path;
66+
/*! \brief The error message if any. */
67+
Optional<String> error_msg;
68+
69+
void VisitAttrs(tvm::AttrVisitor* v) {
70+
v->Visit("artifact_path", &artifact_path);
71+
v->Visit("error_msg", &error_msg);
72+
}
73+
74+
static constexpr const char* _type_key = "meta_schedule.BuilderResult";
75+
TVM_DECLARE_FINAL_OBJECT_INFO(BuilderResultNode, runtime::Object);
76+
};
77+
78+
/*!
79+
* \brief Managed reference to BuilderResultNode
80+
* \sa BuilderResultNode
81+
*/
82+
class BuilderResult : public runtime::ObjectRef {
83+
public:
84+
/*!
85+
* \brief Constructor of BuilderResult.
86+
* \param artifact_path The path to the built artifact.
87+
* \param error_msg The error message if any.
88+
*/
89+
TVM_DLL explicit BuilderResult(Optional<String> artifact_path, Optional<String> error_msg);
90+
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderResult, runtime::ObjectRef, BuilderResultNode);
91+
};
92+
93+
/*! \brief The abstract builder interface. */
94+
class BuilderNode : public runtime::Object {
95+
public:
96+
/*! \brief Default destructor */
97+
virtual ~BuilderNode() = default;
98+
/*!
99+
* \brief Generate the build results from build inputs.
100+
* \param build_inputs The inputs to be built.
101+
* \return The build results.
102+
*/
103+
virtual Array<BuilderResult> Build(const Array<BuilderInput>& build_inputs) = 0;
104+
/*!
105+
* \brief The function type of `Build` method.
106+
* \param build_inputs The inputs to be built.
107+
* \return The build results.
108+
*/
109+
using FBuild = runtime::TypedPackedFunc<Array<BuilderResult>(const Array<BuilderInput>&)>;
110+
111+
static constexpr const char* _type_key = "meta_schedule.Builder";
112+
TVM_DECLARE_BASE_OBJECT_INFO(BuilderNode, runtime::Object);
113+
};
114+
115+
/*!
116+
* \brief Managed reference to BuilderNode
117+
* \sa BuilderNode
118+
*/
119+
class Builder : public runtime::ObjectRef {
120+
public:
121+
/*!
122+
* \brief Create a builder with customized build method on the python-side.
123+
* \param f_build The packed function to the `Build` function..
124+
* \return The Builder created.
125+
*/
126+
static Builder PyBuilder(BuilderNode::FBuild f_build);
127+
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Builder, runtime::ObjectRef, BuilderNode);
128+
};
129+
130+
/*! \brief An abstract builder with customized build method on the python-side. */
131+
class PyBuilderNode : public BuilderNode {
132+
public:
133+
/*! \brief The packed function to the `Build` function. */
134+
FBuild f_build;
135+
136+
void VisitAttrs(tvm::AttrVisitor* v) {
137+
// `f_build` is not visited
138+
}
139+
140+
Array<BuilderResult> Build(const Array<BuilderInput>& build_inputs) final {
141+
return f_build(build_inputs);
142+
}
143+
144+
static constexpr const char* _type_key = "meta_schedule.PyBuilder";
145+
TVM_DECLARE_FINAL_OBJECT_INFO(PyBuilderNode, BuilderNode);
146+
};
147+
148+
} // namespace meta_schedule
149+
} // namespace tvm
150+
151+
#endif // TVM_META_SCHEDULE_BUILDER_H_

‎python/tvm/meta_schedule/__init__.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Package `tvm.meta_schedule`. The meta schedule infrastructure."""
18+
from . import builder

‎python/tvm/meta_schedule/_ffi_api.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""FFI APIs for tvm.meta_schedule"""
18+
from .._ffi import _init_api
19+
20+
_init_api("meta_schedule", __name__) # pylint: disable=protected-access

‎python/tvm/meta_schedule/builder/__init__.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
The tvm.meta_schedule.builder package.
19+
Meta Schedule builders that translate IRModule to runtime.Module,
20+
and then export
21+
"""
22+
from .builder import Builder, BuilderInput, BuilderResult, PyBuilder
23+
from .local_builder import LocalBuilder

‎python/tvm/meta_schedule/builder/builder.py‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Meta Schedule builders that translate IRModule to runtime.Module, and then export"""
18+
from typing import List, Optional
19+
20+
from tvm._ffi import register_object
21+
from tvm.ir import IRModule
22+
from tvm.runtime import Object
23+
from tvm.target import Target
24+
25+
from .. import _ffi_api
26+
27+
28+
@register_object("meta_schedule.BuilderInput")
29+
class BuilderInput(Object):
30+
"""The builder's input.
31+
32+
Parameters
33+
----------
34+
mod : IRModule
35+
The IRModule to be built.
36+
target : Target
37+
The target to be built for.
38+
"""
39+
40+
mod: IRModule
41+
target: Target
42+
43+
def __init__(self, mod: IRModule, target: Target) -> None:
44+
"""Constructor.
45+
46+
Parameters
47+
----------
48+
mod : IRModule
49+
The IRModule to be built.
50+
target : Target
51+
The target to be built for.
52+
"""
53+
self.__init_handle_by_constructor__(
54+
_ffi_api.BuilderInput, # type: ignore # pylint: disable=no-member
55+
mod,
56+
target,
57+
)
58+
59+
60+
@register_object("meta_schedule.BuilderResult")
61+
class BuilderResult(Object):
62+
"""The builder's result.
63+
64+
Parameters
65+
----------
66+
artifact_path : Optional[str]
67+
The path to the artifact.
68+
error_msg : Optional[str]
69+
The error message.
70+
"""
71+
72+
artifact_path: Optional[str]
73+
error_msg: Optional[str]
74+
75+
def __init__(
76+
self,
77+
artifact_path: Optional[str],
78+
error_msg: Optional[str],
79+
) -> None:
80+
"""Constructor.
81+
82+
Parameters
83+
----------
84+
artifact_path : Optional[str]
85+
The path to the artifact.
86+
error_msg : Optional[str]
87+
The error message.
88+
"""
89+
self.__init_handle_by_constructor__(
90+
_ffi_api.BuilderResult, # type: ignore # pylint: disable=no-member
91+
artifact_path,
92+
error_msg,
93+
)
94+
95+
96+
@register_object("meta_schedule.Builder")
97+
class Builder(Object):
98+
"""The abstract builder interface."""
99+
100+
def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
101+
"""Build the given inputs.
102+
103+
Parameters
104+
----------
105+
build_inputs : List[BuilderInput]
106+
The inputs to be built.
107+
Returns
108+
-------
109+
build_results : List[BuilderResult]
110+
The results of building the given inputs.
111+
"""
112+
return _ffi_api.BuilderBuild(self, build_inputs) # type: ignore # pylint: disable=no-member
113+
114+
115+
@register_object("meta_schedule.PyBuilder")
116+
class PyBuilder(Builder):
117+
"""An abstract builder with customized build method on the python-side."""
118+
119+
def __init__(self):
120+
"""Constructor."""
121+
122+
def f_build(build_inputs: List[BuilderInput]) -> List[BuilderResult]:
123+
return self.build(build_inputs)
124+
125+
self.__init_handle_by_constructor__(
126+
_ffi_api.BuilderPyBuilder, # type: ignore # pylint: disable=no-member
127+
f_build,
128+
)
129+
130+
def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
131+
raise NotImplementedError

0 commit comments

Comments
 (0)