comaniac commented on a change in pull request #9061: URL: https://github.com/apache/tvm/pull/9061#discussion_r715935050
########## File path: python/tvm/meta_schedule/database/json_file.py ########## @@ -0,0 +1,61 @@ +# 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 default database that uses a JSON File to store tuning records""" +from tvm._ffi import register_object + +from .. import _ffi_api +from .database import Database + + +@register_object("meta_schedule.JSONFile") +class JSONFile(Database): + """The class of tuning records. Review comment: According to the discussion, this is more like a database implementation, so it might be better to say something like "The JSON-file-based database". ########## File path: python/tvm/meta_schedule/database/database.py ########## @@ -0,0 +1,240 @@ +# 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. +"""Tuning record database""" +from typing import Any, List + +from tvm._ffi import register_object +from tvm.ir.module import IRModule +from tvm.runtime import Object +from tvm.target import Target +from tvm.tir.schedule import Trace + +from .. import _ffi_api +from ..arg_info import ArgInfo +from ..utils import _json_de_tvm + + +@register_object("meta_schedule.Workload") +class Workload(Object): + """A workload, i.e. an IRModule and its structural hash. + + Parameters + ---------- + mod : IRModule + The workload's IRModule + """ + + mod: IRModule + + def __init__(self, mod: IRModule) -> None: + self.__init_handle_by_constructor__( + _ffi_api.Workload, # type: ignore # pylint: disable=no-member + mod, + ) + + def as_json(self) -> Any: + """Export the workload to a JSON string. + + Returns + ------- + json_str : str + The JSON string exported. + """ + return _json_de_tvm(_ffi_api.WorkloadAsJSON(self)) # type: ignore # pylint: disable=no-member + + @staticmethod + def from_json(json_obj: Any) -> "Workload": + """Create a workload from a json object. + + Parameters + ---------- + json_obj : Any + The json object to parse. + + Returns + ------- + tuning_record : TuningRecord + The parsed tuning record. + """ + return _ffi_api.WorkloadFromJSON(json_obj) # type: ignore # pylint: disable=no-member + + +@register_object("meta_schedule.TuningRecord") +class TuningRecord(Object): + """The class of tuning records. + + Parameters + ---------- + trace : tvm.ir.Trace + The trace of the tuning record. + run_secs : List[float] + The run time of the tuning record. + workload : Workload + The workload of the tuning record. + target : Target + The target of the tuning record. + args_info : List[ArgInfo] + The argument information of the tuning record. + """ + + trace: Trace + run_secs: List[float] + workload: Workload + target: Target + args_info: List[ArgInfo] + + def __init__( + self, + trace: Trace, + run_secs: List[float], + workload: Workload, + target: Target, + args_info: List[ArgInfo], + ) -> None: + self.__init_handle_by_constructor__( + _ffi_api.TuningRecord, # type: ignore # pylint: disable=no-member + trace, + run_secs, + workload, + target, + args_info, + ) + + def as_json(self) -> Any: + """Export the tuning record to a JSON string. + + Returns + ------- + json_str : str + The JSON string exported. + """ + return _json_de_tvm(_ffi_api.TuningRecordAsJSON(self)) # type: ignore # pylint: disable=no-member + + @staticmethod + def from_json(json_obj: Any, workload: Workload) -> "TuningRecord": + """Create a tuning record from a json object. + + Parameters + ---------- + json_obj : Any + The json object to parse. + workload : Workload + The workload. + + Returns + ------- + tuning_record : TuningRecord + The parsed tuning record. + """ + return _ffi_api.TuningRecordFromJSON(json_obj, workload) # type: ignore # pylint: disable=no-member + + +@register_object("meta_schedule.Database") +class Database(Object): + """The abstract database interface.""" + + def commit_workload(self, mod: IRModule) -> Workload: + """Look up or add workload to the database if missing. + + Parameters + ---------- + mod : IRModule + The IRModule to be searched for or added. + + Returns + ------- + workload : Workload + The workload corresponding to the given IRModule. + """ + return _ffi_api.DatabaseCommitWorkload(self, mod) # type: ignore # pylint: disable=no-member + + def commit_tuning_record(self, record: TuningRecord) -> None: + """Add a tuning record to the database. Review comment: ```suggestion """Commit a tuning record to the database. ``` Also it might be better to mention what will happen if the same tuning record is committed twice. ########## File path: include/tvm/meta_schedule/database.h ########## @@ -0,0 +1,275 @@ +/* + * 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_DATABASE_H_ +#define TVM_META_SCHEDULE_DATABASE_H_ + +#include <tvm/meta_schedule/arg_info.h> +#include <tvm/target/target.h> +#include <tvm/tir/schedule/trace.h> + +namespace tvm { +namespace meta_schedule { + +/*! \brief A workload, i.e. an IRModule and its structural hash. */ +class WorkloadNode : public runtime::Object { + public: + /*! \brief The type of structural hash */ + using THashCode = size_t; + /*! \brief The workload's IRModule. */ + IRModule mod; + /*! \brief The workload's structural hash. */ + THashCode shash; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("mod", &mod); + // `shash` is not visited because TVM FFI doesn't support uint64_t + } + + static constexpr const char* _type_key = "meta_schedule.Workload"; + TVM_DECLARE_FINAL_OBJECT_INFO(WorkloadNode, runtime::Object); + + /*! + * \brief Export the workload to a JSON string. + * \return An array containing the structural hash and the base64 json string. + */ + ObjectRef AsJSON() const; +}; + +/*! + * \brief Managed reference to WorkloadNode. + * \sa WorkloadNode + */ +class Workload : public runtime::ObjectRef { + public: + using THashCode = WorkloadNode::THashCode; + /*! + * \brief Constructor of Workload. + * \param mod The workload's IRModule. + */ + TVM_DLL explicit Workload(IRModule mod); + /*! + * \brief Constructor of Workload. + * \param mod The workload's IRModule. + * \param shash The workload's structural hash. + */ + TVM_DLL explicit Workload(IRModule mod, THashCode shash); + /*! + * \brief Create a workload from a json object. + * \param json_obj The json object. + * \return The created workload. + */ + TVM_DLL static Workload FromJSON(const ObjectRef& json_obj); + + TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Workload, runtime::ObjectRef, WorkloadNode); +}; + +/*! \brief The hash method for Workload */ +struct WorkloadHash { + size_t operator()(const Workload& a) const { return a->shash; } +}; + +/*! \brief The equality check for Workload */ +struct WorkloadEqual { + bool operator()(const Workload& a, const Workload& b) const { + return a->shash == b->shash && tvm::StructuralEqual()(a->mod, b->mod); + } +}; + +/*! \brief The class of tuning records. */ +class TuningRecordNode : public runtime::Object { + public: + /*! \brief The trace tuned. */ + tir::Trace trace; + /*! \brief The profiling result in seconds. */ + Array<FloatImm> run_secs; + /*! \brief The workload. */ + Workload workload{nullptr}; + /*! \brief The target for tuning. */ + Target target; + /*! \brief The argument information. */ + Array<ArgInfo> args_info; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("trace", &trace); + v->Visit("run_secs", &run_secs); + v->Visit("workload", &workload); + v->Visit("target", &target); + v->Visit("args_info", &args_info); + } + + static constexpr const char* _type_key = "meta_schedule.TuningRecord"; + TVM_DECLARE_FINAL_OBJECT_INFO(TuningRecordNode, runtime::Object); + + /*! + * \brief Export the tuning record to a JSON string. + * \return An array containing the trace, running secs, serialized target, and + * argument information. + */ + ObjectRef AsJSON() const; +}; + +/*! + * \brief The managed reference of TuningRecordNode. + * \sa TuningRecordNode + */ +class TuningRecord : public runtime::ObjectRef { + public: + /*! + \brief Constructor of a tuning record. + \param trace The trace of the tuning record. + \param run_secs The running time of the tuning record. + \param workload The workload of the tuning record. + \param target The target of the tuning record. + \param args_info The argument information of the tuning record. + */ + TVM_DLL explicit TuningRecord(tir::Trace trace, Array<FloatImm> run_secs, Workload workload, + Target target, Array<ArgInfo> args_info); + /*! + * \brief Create a tuning record from a json object. + * \param json_obj The json object. + * \param workload The workload. + * \return The tuning record created. + */ + TVM_DLL static TuningRecord FromJSON(const ObjectRef& json_obj, const Workload& workload); + TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TuningRecord, runtime::ObjectRef, TuningRecordNode); +}; + +/* \brief The abstract interface of database. */ +class DatabaseNode : public runtime::Object { + public: + /*! \brief Default destructor */ + virtual ~DatabaseNode() = default; + /*! + * \brief Look up or add workload to the database if missing. + * \param mod The IRModule to be searched for or added. + * \return The workload corresponding to the given IRModule. + */ + virtual Workload CommitWorkload(const IRModule& mod) = 0; + /*! + * \brief Add a tuning record to the database. + * \param record The tuning record to be added. + */ + virtual void CommitTuningRecord(const TuningRecord& record) = 0; + /*! + * \brief Get the top K tuning records of given workload from the database. + * \param workload The workload to be searched for. + * \param top_k The number of top records to be returned. + * \return An array of top K tuning records for the given workload. + */ + virtual Array<TuningRecord> GetTopK(const Workload& workload, int top_k) = 0; + /*! + * \brief Get the size of the database. + * \return The size of the database. + */ + virtual int64_t Size() = 0; + + static constexpr const char* _type_key = "meta_schedule.Database"; + TVM_DECLARE_BASE_OBJECT_INFO(DatabaseNode, runtime::Object); +}; + +/*! \brief The database with customized methods on the python-side. */ +class PyDatabaseNode : public DatabaseNode { + public: + /*! + * \brief The function type of `CommitWorkload` method. + * \param mod The IRModule to be searched for or added. + * \return The workload corresponding to the given IRModule. + */ + using FCommitWorkload = runtime::TypedPackedFunc<Workload(const IRModule&)>; + /*! + * \brief The function type of `CommitTuningRecord` method. + * \param record The tuning record to be added. + */ + using FCommitTuningRecord = runtime::TypedPackedFunc<void(const TuningRecord&)>; + /*! + * \brief The function type of `GetTopK` method. + * \param workload The workload to be searched for. + * \param top_k The number of top records to be returned. + * \return An array of top K tuning records for the given workload. + */ + using FGetTopK = runtime::TypedPackedFunc<Array<TuningRecord>(const Workload&, int)>; + /*! + * \brief The function type of `Size` method. + * \return The size of the database. + */ + using FSize = runtime::TypedPackedFunc<int64_t()>; + + /*! \brief The packed function to the `CommitWorkload` function. */ + FCommitWorkload f_commit_workload; + /*! \brief The packed function to the `CommitTuningRecord` function. */ + FCommitTuningRecord f_commit_tuning_record; + /*! \brief The packed function to the `GetTopK` function. */ + FGetTopK f_get_top_k; + /*! \brief The packed function to the `Size` function. */ + FSize f_size; + + void VisitAttrs(tvm::AttrVisitor* v) { + // `f_commit_workload` is not visited + // `f_commit_tuning_record` is not visited + // `f_get_top_k` is not visited + // `f_size` is not visited + } + + static constexpr const char* _type_key = "meta_schedule.PyDatabase"; + TVM_DECLARE_FINAL_OBJECT_INFO(PyDatabaseNode, DatabaseNode); + + Workload CommitWorkload(const IRModule& mod) final { return f_commit_workload(mod); } + + void CommitTuningRecord(const TuningRecord& record) final { f_commit_tuning_record(record); } + + Array<TuningRecord> GetTopK(const Workload& workload, int top_k) final { + return f_get_top_k(workload, top_k); + } + + int64_t Size() final { return f_size(); } +}; + +/*! + * \brief Managed reference to DatabaseNode. + * \sa DatabaseNode + */ +class Database : public runtime::ObjectRef { + public: + /*! + * \brief Create a default database that uses JSON file for tuning records. + * \param path_workload The path to the workload table. + * \param path_tuning_record The path to the database table. + * \param allow_missing Whether to create new file when the given path is not found. + */ + TVM_DLL static Database JSONFile(String path_workload, String path_tuning_record, + bool allow_missing); Review comment: After the refactoring the logic is clearer to me, and I also realized that this API is a bit confusing. Per offline discussion with @junrushao1994, this is more like one of the factory methods. In the future, we might have: ``` class Database { static Database JSONFile(...); static Database SQLite(...); }; ``` And uses might be able to initialize a DB for their needs: `Database::SQLite(...)`. Accordingly, it might be better to consistently name them as database names. For example, we could name `JSONFile` to be `JsonDB`, but I'm bad at naming so would be happy to have better ideas from others. cc @mbrookhart ########## File path: python/tvm/meta_schedule/database/database.py ########## @@ -0,0 +1,240 @@ +# 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. +"""Tuning record database""" +from typing import Any, List + +from tvm._ffi import register_object +from tvm.ir.module import IRModule +from tvm.runtime import Object +from tvm.target import Target +from tvm.tir.schedule import Trace + +from .. import _ffi_api +from ..arg_info import ArgInfo +from ..utils import _json_de_tvm + + +@register_object("meta_schedule.Workload") +class Workload(Object): + """A workload, i.e. an IRModule and its structural hash. + + Parameters + ---------- + mod : IRModule + The workload's IRModule + """ + + mod: IRModule + + def __init__(self, mod: IRModule) -> None: + self.__init_handle_by_constructor__( + _ffi_api.Workload, # type: ignore # pylint: disable=no-member + mod, + ) + + def as_json(self) -> Any: + """Export the workload to a JSON string. + + Returns + ------- + json_str : str + The JSON string exported. + """ + return _json_de_tvm(_ffi_api.WorkloadAsJSON(self)) # type: ignore # pylint: disable=no-member + + @staticmethod + def from_json(json_obj: Any) -> "Workload": + """Create a workload from a json object. + + Parameters + ---------- + json_obj : Any + The json object to parse. + + Returns + ------- + tuning_record : TuningRecord + The parsed tuning record. + """ + return _ffi_api.WorkloadFromJSON(json_obj) # type: ignore # pylint: disable=no-member + + +@register_object("meta_schedule.TuningRecord") +class TuningRecord(Object): + """The class of tuning records. + + Parameters + ---------- + trace : tvm.ir.Trace + The trace of the tuning record. + run_secs : List[float] + The run time of the tuning record. + workload : Workload + The workload of the tuning record. + target : Target + The target of the tuning record. + args_info : List[ArgInfo] + The argument information of the tuning record. + """ + + trace: Trace + run_secs: List[float] + workload: Workload + target: Target + args_info: List[ArgInfo] + + def __init__( + self, + trace: Trace, + run_secs: List[float], + workload: Workload, + target: Target, + args_info: List[ArgInfo], + ) -> None: + self.__init_handle_by_constructor__( + _ffi_api.TuningRecord, # type: ignore # pylint: disable=no-member + trace, + run_secs, + workload, + target, + args_info, + ) + + def as_json(self) -> Any: + """Export the tuning record to a JSON string. + + Returns + ------- + json_str : str + The JSON string exported. + """ + return _json_de_tvm(_ffi_api.TuningRecordAsJSON(self)) # type: ignore # pylint: disable=no-member + + @staticmethod + def from_json(json_obj: Any, workload: Workload) -> "TuningRecord": + """Create a tuning record from a json object. + + Parameters + ---------- + json_obj : Any + The json object to parse. + workload : Workload + The workload. + + Returns + ------- + tuning_record : TuningRecord + The parsed tuning record. + """ + return _ffi_api.TuningRecordFromJSON(json_obj, workload) # type: ignore # pylint: disable=no-member + + +@register_object("meta_schedule.Database") +class Database(Object): + """The abstract database interface.""" + + def commit_workload(self, mod: IRModule) -> Workload: + """Look up or add workload to the database if missing. Review comment: ```suggestion """Commit a workload to the database if missing. ``` ########## File path: tests/python/unittest/test_meta_schedule_database.py ########## @@ -0,0 +1,274 @@ +# 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. +# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring +"""Test Meta Schedule Database""" +import os.path as osp +import sys +import tempfile +from typing import Callable + +import pytest + +import tvm +from tvm import tir +from tvm.ir.module import IRModule +from tvm.meta_schedule.arg_info import ArgInfo +from tvm.meta_schedule.database import JSONFile, TuningRecord +from tvm.script import ty +from tvm.tir import Schedule + +# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument +# fmt: off + [email protected] +class Matmul: + def main(a: ty.handle, b: ty.handle, c: ty.handle) -> None: + tir.func_attr({"global_symbol": "main"}) + A = tir.match_buffer(a, (1024, 1024), "float32") + B = tir.match_buffer(b, (1024, 1024), "float32") + C = tir.match_buffer(c, (1024, 1024), "float32") + with tir.block([1024, 1024, tir.reduce_axis(0, 1024)], "matmul") as [vi, vj, vk]: + with tir.init(): + C[vi, vj] = 0.0 + C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj] + + [email protected] +class MatmulRelu: + def main(a: ty.handle, b: ty.handle, d: ty.handle) -> None: # pylint: disable=no-self-argument + tir.func_attr({"global_symbol": "main", "tir.noalias": True}) + A = tir.match_buffer(a, (16, 16), "float32") + B = tir.match_buffer(b, (16, 16), "float32") + D = tir.match_buffer(d, (16, 16), "float32") + C = tir.alloc_buffer((16, 16), "float32") + with tir.block([16, 16, tir.reduce_axis(0, 16)], "matmul") as [vi, vj, vk]: + with tir.init(): + C[vi, vj] = 0.0 + C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj] + with tir.block([16, 16], "relu") as [vi, vj]: + D[vi, vj] = tir.max(C[vi, vj], 0.0) + + +# fmt: on +# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument + + +def _schedule_matmul(sch: Schedule): + block = sch.get_block("matmul") + i, j, k = sch.get_loops(block=block) + i_tiles = [1, 1, 2, 512] + j_tiles = [1, 512, 1, 2] + k_tiles = [256, 4] + i_0, i_1, i_2, i_3 = sch.split(loop=i, factors=i_tiles) + j_0, j_1, j_2, j_3 = sch.split(loop=j, factors=j_tiles) + k_0, k_1 = sch.split(loop=k, factors=k_tiles) + sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3) + + +def _create_schedule(mod: IRModule, sch_fn: Callable[[Schedule], None]) -> Schedule: + sch = tir.Schedule(mod=mod, debug_mask="all") + sch_fn(sch) + return sch + + +def _create_tmp_database(tmpdir: str) -> JSONFile: + path_workload = osp.join(tmpdir, "workloads.json") + path_tuning_record = osp.join(tmpdir, "tuning_records.json") + return JSONFile(path_workload, path_tuning_record) Review comment: Echo again. It seems much more straightforward to write `return JsonDB(path_workload, path_tuning_record)` here. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
