Merge branch 'main' into aluo/onnx/unsqueeze-alt

* main:
  Add back missing __init__.py to unbreak CI. (apache#9052)
  [Meta Schedule][M3b] Builder (apache#9044)

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