[MetaSchedule][M3b] Builder

CMakeLists.txt

@@ -240,6 +240,7 @@ assign_source_group("Include" ${GROUP_INCLUDE})
 # Source file lists
 file(GLOB_RECURSE COMPILER_SRCS
  src/auto_scheduler/*.cc
+ src/meta_schedule/*.cc
  src/node/*.cc
  src/ir/*.cc
  src/arith/*.cc

include/tvm/meta_schedule/builder.h

@@ -0,0 +1,151 @@
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_META_SCHEDULE_BUILDER_H_
+#define TVM_META_SCHEDULE_BUILDER_H_
+
+#include <tvm/ir/module.h>
+#include <tvm/target/target.h>
+
+namespace tvm {
+namespace meta_schedule {
+
+/*! \brief The builder's input. */
+class BuilderInputNode : public runtime::Object {
+ public:
+ /*! \brief The IRModule to be built. */
+ IRModule mod;
+ /*! \brief The target to be built for. */
+ Target target;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("mod", &mod);
+ v->Visit("target", &target);
+ }
+
+ static constexpr const char* _type_key = "meta_schedule.BuilderInput";
+ TVM_DECLARE_FINAL_OBJECT_INFO(BuilderInputNode, runtime::Object);
+};
+
+/*!
+ * \brief Managed reference to BuilderInputNode
+ * \sa BuilderInputNode
+ */
+class BuilderInput : public runtime::ObjectRef {
+ public:
+ /*!
+ * \brief Constructor of BuilderInput.
+ * \param mod The IRModule to be built.
+ * \param target The target to be built for.
+ */
+ TVM_DLL explicit BuilderInput(IRModule mod, Target target);
+ TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderInput, runtime::ObjectRef, BuilderInputNode);
+};
+
+/*! \brief The builder's output. */
+class BuilderResultNode : public runtime::Object {
+ public:
+ /*! \brief The path to the built artifact. */
+ Optional<String> artifact_path;
+ /*! \brief The error message if any. */
+ Optional<String> error_msg;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("artifact_path", &artifact_path);
+ v->Visit("error_msg", &error_msg);
+ }
+
+ static constexpr const char* _type_key = "meta_schedule.BuilderResult";
+ TVM_DECLARE_FINAL_OBJECT_INFO(BuilderResultNode, runtime::Object);
+};
+
+/*!
+ * \brief Managed reference to BuilderResultNode
+ * \sa BuilderResultNode
+ */
+class BuilderResult : public runtime::ObjectRef {
+ public:
+ /*!
+ * \brief Constructor of BuilderResult.
+ * \param artifact_path The path to the built artifact.
+ * \param error_msg The error message if any.
+ */
+ TVM_DLL explicit BuilderResult(Optional<String> artifact_path, Optional<String> error_msg);
+ TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(BuilderResult, runtime::ObjectRef, BuilderResultNode);
+};
+
+/*! \brief The abstract builder interface. */
+class BuilderNode : public runtime::Object {
+ public:
+ /*! \brief Default destructor */
+ virtual ~BuilderNode() = default;
+ /*!
+ * \brief Generate the build results from build inputs.
+ * \param build_inputs The inputs to be built.
+ * \return The build results.
+ */
+ virtual Array<BuilderResult> Build(const Array<BuilderInput>& build_inputs) = 0;
+ /*!
+ * \brief The function type of `Build` method.
+ * \param build_inputs The inputs to be built.
+ * \return The build results.
+ */
+ using FBuild = runtime::TypedPackedFunc<Array<BuilderResult>(const Array<BuilderInput>&)>;
+
+ static constexpr const char* _type_key = "meta_schedule.Builder";
+ TVM_DECLARE_BASE_OBJECT_INFO(BuilderNode, runtime::Object);
+};
+
+/*!
+ * \brief Managed reference to BuilderNode
+ * \sa BuilderNode
+ */
+class Builder : public runtime::ObjectRef {
+ public:
+ /*!
+ * \brief Create a builder with customized build method on the python-side.
+ * \param f_build The packed function to the `Build` function..
+ * \return The Builder created.
+ */
+ static Builder PyBuilder(BuilderNode::FBuild f_build);
+ TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Builder, runtime::ObjectRef, BuilderNode);
+};
+
+/*! \brief An abstract builder with customized build method on the python-side. */
+class PyBuilderNode : public BuilderNode {
+ public:
+ /*! \brief The packed function to the `Build` function. */
+ FBuild f_build;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ // `f_build` is not visited
+ }
+
+ Array<BuilderResult> Build(const Array<BuilderInput>& build_inputs) final {
+ return f_build(build_inputs);
+ }
+
+ static constexpr const char* _type_key = "meta_schedule.PyBuilder";
+ TVM_DECLARE_FINAL_OBJECT_INFO(PyBuilderNode, BuilderNode);
+};
+
+} // namespace meta_schedule
+} // namespace tvm
+
+#endif // TVM_META_SCHEDULE_BUILDER_H_

python/tvm/meta_schedule/__init__.py

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

python/tvm/meta_schedule/_ffi_api.py

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

python/tvm/meta_schedule/builder/__init__.py

@@ -0,0 +1,23 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The tvm.meta_schedule.builder package.
+Meta Schedule builders that translate IRModule to runtime.Module,
+and then export
+"""
+from .builder import Builder, BuilderInput, BuilderResult, PyBuilder
+from .local_builder import LocalBuilder

python/tvm/meta_schedule/builder/builder.py

@@ -0,0 +1,131 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Meta Schedule builders that translate IRModule to runtime.Module, and then export"""
+from typing import List, Optional
+
+from tvm._ffi import register_object
+from tvm.ir import IRModule
+from tvm.runtime import Object
+from tvm.target import Target
+
+from .. import _ffi_api
+
+
+@register_object("meta_schedule.BuilderInput")
+class BuilderInput(Object):
+ """The builder's input.
+
+ Parameters
+ ----------
+ mod : IRModule
+ The IRModule to be built.
+ target : Target
+ The target to be built for.
+ """
+
+ mod: IRModule
+ target: Target
+
+ def __init__(self, mod: IRModule, target: Target) -> None:
+ """Constructor.
+
+ Parameters
+ ----------
+ mod : IRModule
+ The IRModule to be built.
+ target : Target
+ The target to be built for.
+ """
+ self.__init_handle_by_constructor__(
+ _ffi_api.BuilderInput, # type: ignore # pylint: disable=no-member
+ mod,
+ target,
+ )
+
+
+@register_object("meta_schedule.BuilderResult")
+class BuilderResult(Object):
+ """The builder's result.
+
+ Parameters
+ ----------
+ artifact_path : Optional[str]
+ The path to the artifact.
+ error_msg : Optional[str]
+ The error message.
+ """
+
+ artifact_path: Optional[str]
+ error_msg: Optional[str]
+
+ def __init__(
+ self,
+ artifact_path: Optional[str],
+ error_msg: Optional[str],
+ ) -> None:
+ """Constructor.
+
+ Parameters
+ ----------
+ artifact_path : Optional[str]
+ The path to the artifact.
+ error_msg : Optional[str]
+ The error message.
+ """
+ self.__init_handle_by_constructor__(
+ _ffi_api.BuilderResult, # type: ignore # pylint: disable=no-member
+ artifact_path,
+ error_msg,
+ )
+
+
+@register_object("meta_schedule.Builder")
+class Builder(Object):
+ """The abstract builder interface."""
+
+ def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
+ """Build the given inputs.
+
+ Parameters
+ ----------
+ build_inputs : List[BuilderInput]
+ The inputs to be built.
+ Returns
+ -------
+ build_results : List[BuilderResult]
+ The results of building the given inputs.
+ """
+ return _ffi_api.BuilderBuild(self, build_inputs) # type: ignore # pylint: disable=no-member
+
+
+@register_object("meta_schedule.PyBuilder")
+class PyBuilder(Builder):
+ """An abstract builder with customized build method on the python-side."""
+
+ def __init__(self):
+ """Constructor."""
+
+ def f_build(build_inputs: List[BuilderInput]) -> List[BuilderResult]:
+ return self.build(build_inputs)
+
+ self.__init_handle_by_constructor__(
+ _ffi_api.BuilderPyBuilder, # type: ignore # pylint: disable=no-member
+ f_build,
+ )
+
+ def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
+ raise NotImplementedError