You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/09/25 05:48:26 UTC
[GitHub] [tvm] junrushao opened a new pull request, #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
junrushao opened a new pull request, #12895:
URL: https://github.com/apache/tvm/pull/12895
This PR intends to refactor the existing tuning API into a new paradigm to enhance developer ergonomics.
The original behavior assumes tuning is an end-to-end process that transforms an IRModule into a runtime Module.
```python
def tune_relay(...) -> runtime.Module: ...
```
This PR refactors this behavior into two stages, i.e. tuning and compilation, where tuning is time-consuming, while compilation is supposed to be swift.
```python
def tune_relay(...) -> Database: ...
def compile_relay_with_database(...) -> runtime.Module: ...
```
The PR is intended to be open for at least one week as it's a significant change.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r984247929
##########
include/tvm/meta_schedule/search_strategy.h:
##########
@@ -257,8 +249,8 @@ class PySearchStrategyNode : public SearchStrategyNode {
}
void InitializeWithTuneContext(const TuneContext& context) final;
- void PreTuning(const Array<tir::Schedule>& design_spaces, const Optional<Database>& database,
- const Optional<CostModel>& cost_model) final;
+ void PreTuning(int max_trials, int num_trials_per_iter, const Array<tir::Schedule>& design_spaces,
Review Comment:
May I ask how `max_trials_per_task` is used in the new interface?
##########
include/tvm/meta_schedule/task_scheduler.h:
##########
@@ -143,55 +202,43 @@ class TaskScheduler;
/*! \brief The task scheduler with customized methods on the python-side. */
class PyTaskSchedulerNode : public TaskSchedulerNode {
public:
- /*! \brief The function type of `Tune` method. */
- using FTune = runtime::TypedPackedFunc<void()>;
-
- /*! \brief The function type of `InitializeTask` method. */
- using FInitializeTask = runtime::TypedPackedFunc<void(int)>;
-
/*!
- * \brief The function type of `TouchTask` method.
- * \param task_id The task id to be checked.
- * \return Whether the task is running.
+ * \brief The function type of `NextTaskId` method.
+ * \return The next task id.
*/
- using FTouchTask = runtime::TypedPackedFunc<void(int)>;
-
+ using FNextTaskId = runtime::TypedPackedFunc<int()>;
/*!
* \brief The function type of `JoinRunningTask` method.
* \param task_id The task id to be joined.
*/
using FJoinRunningTask = runtime::TypedPackedFunc<Array<RunnerResult>(int)>;
+ /*! \brief The function type of `Tune` method. */
+ using FTune = runtime::TypedPackedFunc<void(
Review Comment:
Ditto. Format.
##########
python/tvm/meta_schedule/cost_model/cost_model.py:
##########
@@ -16,23 +16,30 @@
# under the License.
"""Meta Schedule CostModel."""
import ctypes
-from typing import Callable, List
+from typing import Callable, List, Union
+
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm._ffi import register_object
-from tvm.meta_schedule.utils import _get_default_str
from tvm.runtime import Object
from .. import _ffi_api
from ..runner import RunnerResult
from ..search_strategy import MeasureCandidate
from ..tune_context import TuneContext
+from ..utils import _get_default_str
@register_object("meta_schedule.CostModel")
class CostModel(Object):
"""Cost model."""
+ CostModelType = Union["CostModel", Literal["xgb", "mlp", "random"]]
Review Comment:
Maybe we need a note/comment somewhere to remind developer to update the list here when some new cost model is created.
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,82 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+# isort: off
Review Comment:
This could be put behind the import if there is only one import that we don't want to sort. Btw why do we put isort here for this import?
##########
include/tvm/support/random_engine.h:
##########
@@ -46,32 +43,18 @@ namespace support {
class LinearCongruentialEngine {
public:
- /*!
- * \brief The result type is defined as uint64_t here to avoid overflow.
- * \note The type name is not in Google style because it is used in STL's distribution inferface.
Review Comment:
Suggest to keep for better code understanding.
##########
include/tvm/meta_schedule/task_scheduler.h:
##########
@@ -73,66 +127,71 @@ namespace meta_schedule {
*/
class TaskSchedulerNode : public runtime::Object {
public:
- /*! \brief The tasks to be tuned */
- Array<TuneContext> tasks;
- /*! \brief The builder of the scheduler. */
- Builder builder{nullptr};
- /*! \brief The runner of the scheduler. */
- Runner runner{nullptr};
- /*! \brief The database of the scheduler. */
- Optional<Database> database;
- /*! \brief The cost model of the scheduler. */
- Optional<CostModel> cost_model;
+ /*! \brief The tuning task's logging function. */
+ PackedFunc logger;
+ /*! \brief Records for each task */
+ Array<TaskRecord> tasks_;
/*! \brief The list of measure callbacks of the scheduler. */
- Array<MeasureCallback> measure_callbacks;
- /*! \brief The maximum number of trials allowed. */
- int max_trials;
- /*! \brief The number of trials already conducted. */
- int num_trials_already;
- /*! \brief The tuning task's logging function. t*/
- PackedFunc logging_func;
+ Array<MeasureCallback> measure_callbacks_;
+ /*! \brief The database used in tuning */
+ Optional<Database> database_;
+ /*! \brief The cost model used in tuning */
+ Optional<CostModel> cost_model_;
+ /*! \brief The number of remaining tasks to be tuned. */
+ int remaining_tasks_;
/*! \brief The default destructor. */
virtual ~TaskSchedulerNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) {
- v->Visit("tasks", &tasks);
- v->Visit("builder", &builder);
- v->Visit("runner", &runner);
- v->Visit("database", &database);
- v->Visit("cost_model", &cost_model);
- v->Visit("measure_callbacks", &measure_callbacks);
- v->Visit("max_trials", &max_trials);
- v->Visit("num_trials_already", &num_trials_already);
- // `logging_func` is not visited
+ // `logger` is not visited
+ v->Visit("tasks_", &tasks_);
+ v->Visit("measure_callbacks_", &measure_callbacks_);
+ v->Visit("database_", &database_);
+ v->Visit("cost_model_", &cost_model_);
+ v->Visit("remaining_tasks_", &remaining_tasks_);
}
- /*! \brief Auto-tuning. */
- virtual void Tune();
-
- /*!
- * \brief Initialize modules of the given task.
- * \param task_id The task id to be initialized.
- */
- virtual void InitializeTask(int task_id);
-
/*!
- * \brief Touch the task and update its status
- * \param task_id The task id to be checked.
+ * \brief Fetch the next task id.
+ * \return The next task id.
*/
- virtual void TouchTask(int task_id);
-
+ virtual int NextTaskId() = 0;
/*!
* \brief Wait until the task is finished.
* \param task_id The task id to be joined.
+ * \return The results from the runner.
*/
virtual Array<RunnerResult> JoinRunningTask(int task_id);
-
/*!
- * \brief Fetch the next task id.
- * \return The next task id.
+ * \brief Jointly tune a given list of tasks.
+ * \param tasks The tasks to be tuned
+ * \param task_weights The weight of each task
+ * \param max_trials_global The maximum number of trials to be performed globally
+ * \param max_trials_per_task The maximum number of trials to be performed for each task
+ * \param num_trials_per_iter The number of trials to be performed in each iteration
+ * \param builder The MetaSchedule builder
+ * \param runner The MetaSchedule runner
+ * \param measure_callbacks The callbacks to be called after each measurement
+ * \param database The database used in tuning
+ * \param cost_model The cost model used in tuning
*/
- virtual int NextTaskId() = 0;
+ virtual void Tune(Array<TuneContext> tasks, Array<FloatImm> task_weights, int max_trials_global,
Review Comment:
Can you line these arguments up for better code formatting?
##########
python/tvm/meta_schedule/logging.py:
##########
@@ -0,0 +1,262 @@
+# 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.
+"""Logging interface in MetaSchedule"""
+import logging
+import logging.config
+import os
+import os.path as osp
+from logging import Logger
+from typing import Any, Callable, Dict, List, Optional
+
+
+def get_logger(name: str) -> Logger:
+ """Create or get a logger by its name. This is essentially a wrapper of python's native logger.
+
+ Parameters
+ ----------
+ name : str
+ The name of the logger.
+
+ Returns
+ -------
+ logger : Logger
+ The logger instance.
+ """
+ return logging.getLogger(name)
+
+
+def make_logging_func(logger: Logger) -> Optional[Callable[[int, str], None]]:
Review Comment:
`get_logging_func` sounds more reasonable.
##########
include/tvm/meta_schedule/space_generator.h:
##########
@@ -171,6 +213,7 @@ class PySpaceGeneratorNode : public SpaceGeneratorNode {
FClone f_clone;
void VisitAttrs(tvm::AttrVisitor* v) {
+ SpaceGeneratorNode::VisitAttrs(v);
Review Comment:
Is this expected for all derived classes?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r982961786
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}},
Review Comment:
oh thanks! sorry for being sloppy copying things around. fixed. would you like to take another look? thanks a bunch!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988321661
##########
python/tvm/contrib/torch/optimize_torch.py:
##########
@@ -62,75 +51,32 @@ def forward(self, *torch_inputs: Tuple[torch.Tensor]):
return ret
-def llvm_target():
- return "llvm -num-cores"
-
-
@register_func("script_torch.save_to_base64")
def save_to_base64(obj) -> bytes:
with tempfile.NamedTemporaryFile(suffix=".so") as tmpfile:
obj.export_library(tmpfile.name)
- with open(tmpfile.name, "rb") as tfile:
- return base64.b64encode(tfile.read())
-
-
-def tune_relay_auto(
- mod: IRModule,
- target: Union[str, Target],
- config: TuneConfig,
- work_dir: str,
- backend: str = "graph",
- params: Optional[Dict[str, NDArray]] = None,
-) -> Union[Module, vm.Executable]:
- """A wrapper of `tune_relay` but provide a default setting for the config.
-
- Parameters
- ----------
- mod : IRModule
- The module to tune.
- target : Union[str, Target]
- The target to tune for.
- config : TuneConfig
- The search strategy config.
- params : Optional[Dict[str, tvm.runtime.NDArray]]
- The associated parameters of the program
- work_dir : Optional[str]
- The working directory to save intermediate results.
- backend : str = "graph"
- The backend to use for relay compilation(graph / vm).
-
- Returns
- -------
- lib : Union[Module, tvm.runtime.vm.Executable]
- The built runtime module or vm Executable for the given relay workload.
- """
- target = default_config.target(target)
- extracted_tasks = extract_task_from_relay(mod, target, params)
- if config is None:
- config = TuneConfig(
- num_trials_per_iter=16,
- max_trials_global=16 * len(extracted_tasks),
- )
- database = tune_extracted_tasks(extracted_tasks, config, work_dir)
- relay_build = {"graph": relay.build, "vm": relay.vm.compile}[backend]
- with target, autotvm_silencer(), database:
- with PassContext(
- opt_level=3,
- config={
- "relay.backend.use_meta_schedule": True,
- "relay.backend.use_meta_schedule_dispatch": target.kind.name != "cuda",
- "relay.backend.tir_converter": "default",
- },
- ):
- return relay_build(mod, target=target, params=params)
+ with open(tmpfile.name, "rb") as temp_file:
+ return base64.b64encode(temp_file.read())
def optimize_torch(
func,
example_inputs,
- tuning_config=None,
- target=None,
+ *,
work_dir=None,
+ target: Union[str, Target] = "cpu",
+ max_trials_global: int,
+ max_trials_per_task: Optional[int] = None,
+ num_trials_per_iter: int = 64,
+ builder: ms.Builder.BuilderType = "local",
+ runner: ms.Runner.RunnerType = "local",
+ database: ms.Database.DatabaseType = "json",
+ cost_model: ms.CostModel.CostModelType = "xgb",
+ measure_callbacks: ms.MeasureCallback.CallbackListType = "default",
+ task_scheduler: ms.TaskScheduler.TaskSchedulerType = "gradient",
+ space: ms.SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+ strategy: ms.SearchStrategy.SearchStrategyType = "evolutionary",
Review Comment:
added
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190735
##########
include/tvm/meta_schedule/task_scheduler.h:
##########
@@ -143,55 +202,43 @@ class TaskScheduler;
/*! \brief The task scheduler with customized methods on the python-side. */
class PyTaskSchedulerNode : public TaskSchedulerNode {
public:
- /*! \brief The function type of `Tune` method. */
- using FTune = runtime::TypedPackedFunc<void()>;
-
- /*! \brief The function type of `InitializeTask` method. */
- using FInitializeTask = runtime::TypedPackedFunc<void(int)>;
-
/*!
- * \brief The function type of `TouchTask` method.
- * \param task_id The task id to be checked.
- * \return Whether the task is running.
+ * \brief The function type of `NextTaskId` method.
+ * \return The next task id.
*/
- using FTouchTask = runtime::TypedPackedFunc<void(int)>;
-
+ using FNextTaskId = runtime::TypedPackedFunc<int()>;
/*!
* \brief The function type of `JoinRunningTask` method.
* \param task_id The task id to be joined.
*/
using FJoinRunningTask = runtime::TypedPackedFunc<Array<RunnerResult>(int)>;
+ /*! \brief The function type of `Tune` method. */
+ using FTune = runtime::TypedPackedFunc<void(
Review Comment:
done
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989280921
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,101 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+if TYPE_CHECKING:
+ from tvm import relay
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=True,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+ pass_config: Mapping[str, Any],
+ executor: Optional["relay.backend.Executor"],
+) -> Tuple[IRModule, Target, Dict[str, NDArray], Dict[str, Any], "relay.backend.Executor"]:
+ from tvm import relay # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, relay.Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+ if executor is None:
+ executor = relay.backend.Executor("graph")
+ pass_config = dict(pass_config)
+ pass_config.setdefault(
+ "relay.FuseOps.link_params",
+ executor.attrs.get("link_params", False),
+ )
Review Comment:
so the idea is that we never use `relay.FuseOps.link_params`, but instead attaching executor to the IRModule - is my understanding correct?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1264223324
@sunggg would you also give a review to confirm if there isn’t any potential issue with Relax integration?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r983969770
##########
python/tvm/meta_schedule/logging.py:
##########
@@ -0,0 +1,262 @@
+# 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.
+"""Logging interface in MetaSchedule"""
+import logging
+import logging.config
+import os
+import os.path as osp
+from logging import Logger
+from typing import Any, Callable, Dict, List, Optional
+
+
+def get_logger(name: str) -> Logger:
+ """Create or get a logger by its name. This is essentially a wrapper of python's native logger.
+
+ Parameters
+ ----------
+ name : str
+ The name of the logger.
+
+ Returns
+ -------
+ logger : Logger
+ The logger instance.
+ """
+ return logging.getLogger(name)
+
+
+def make_logging_func(logger: Logger) -> Optional[Callable[[int, str], None]]:
+ """Get the logging function.
+
+ Parameters
+ ----------
+ logger : Logger
+ The logger instance.
+ Returns
+ -------
+ result : Optional[Callable]
+ The function to do the specified level of logging.
+ """
+ if logger is None:
+ return None
+
+ level2log = {
+ logging.DEBUG: logger.debug,
+ logging.INFO: logger.info,
+ logging.WARNING: logger.warning,
+ logging.ERROR: logger.error,
+ # logging.FATAL not included
+ }
+
+ def logging_func(level: int, msg: str):
+ def clear_notebook_output():
+ from IPython.display import ( # type: ignore # pylint: disable=import-outside-toplevel
+ clear_output,
+ )
+
+ clear_output(wait=True)
+
+ if level < 0:
+ clear_notebook_output()
Review Comment:
Yes, I defined `level < 0` to clear logging for console tuning as well.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985051150
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
Review Comment:
Thanks for spotting this!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1269261207
Hey @masahi, I added `executor` parameter to `extract_tasks` and `compile_relay`, which controls the default value of `relay.FuseOps.link_params` in pass configuration. It's quite confusing to me that `executor` is lifted out of pass config and somehow control the compilation process in a half-functioning way (only works for `GraphExecutor`), and am not sure if I'm using that correctly, so please feel free to suggest what the best way is :-)
On the other hand, as a high-level API, I would prefer not to tweak `tune_relay` adding more parameters to it, given we wanted to give a cleaner interface to introductory level users. Instead, advance users could always use `extract_tasks` + `tune_tasks` + `compile_relay` to get fine-grained control over the tuning process
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r982931215
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}},
Review Comment:
ooop. would you mind pointing out the correct one? thanks a lot!
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}},
Review Comment:
ooops! would you mind pointing out the correct one? thanks a lot!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r983980554
##########
python/tvm/meta_schedule/utils.py:
##########
@@ -114,7 +113,10 @@ def __init__(self, *args, **kwargs):
def __getattr__(self, name: str):
"""Bridge the attribute function."""
- return self._inst.__getattribute__(name)
+ try:
+ return self._inst.__getattribute__(name)
+ except AttributeError:
+ return super(TVMDerivedObject, self).__getattr__(name)
Review Comment:
I don't have much idea how to fix the false positive warnings from mypy. Any suggestions?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988422820
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -69,47 +124,229 @@ def extract_task_from_relay(
"""
# pylint: disable=import-outside-toplevel
from tvm import autotvm
- from tvm.relay import Function as RelayFunc
# pylint: enable=import-outside-toplevel
+ mod, target, params = _normalize_params(mod, target, params)
+ pass_config = dict(pass_config)
+ if target.kind.name != "cuda" and isinstance(
+ autotvm.DispatchContext.current, autotvm.FallbackContext
+ ):
+ tophub_context = autotvm.tophub.context(target)
+ else:
+ tophub_context = autotvm.utils.EmptyContext()
+ with Profiler.timeit("TaskExtraction"):
+ with target, _autotvm_silencer(), tophub_context:
+ with transform.PassContext(
+ opt_level=opt_level,
+ config=pass_config,
+ ):
+ return list(_extract_task(mod, target, params))
+
+
+def extracted_tasks_to_tune_contexts(
+ extracted_tasks: List[ExtractedTask],
+ work_dir: str,
+ space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+ strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+ num_threads: Union[Literal["physical", "logical"], int] = "physical",
+ seed: Optional[int] = None,
+) -> Tuple[List[TuneContext], List[float]]:
+ """Convert ExtractedTask to TuneContext.
+
+ Parameters
+ ----------
+ tasks : List[ExtractedTask]
+ The tasks to be converted
+ work_dir : str
+ The working directory to store logs and databases
+ space : SpaceGenerator.SpaceGeneratorType
+ The space generator to use.
+ strategy : SearchStrategy.SearchStrategyType
+ The search strategy to use.
+ num_threads : Union[Literal["physical", "logical"], int]
+ The number of threads to use.
+ seed : Optional[int]
+ The random seed to use.
+
+ Returns
+ -------
+ tasks : List[TuneContext]
+ The converted tasks
+ task_weights : List[float]
+ The weights of the tasks
+ """
+ tasks: List[TuneContext] = []
+ task_weights: List[float] = []
+ for task, logger, rand_state in zip(
+ extracted_tasks,
+ get_loggers_from_work_dir(work_dir, [t.task_name for t in extracted_tasks]),
+ fork_seed(seed, n=len(extracted_tasks)),
+ ):
+ tasks.append(
+ TuneContext(
+ mod=task.dispatched[0],
+ target=task.target,
+ space_generator=space,
+ search_strategy=strategy,
+ task_name=task.task_name,
+ logger=logger,
+ rand_state=rand_state,
+ num_threads=num_threads,
+ ).clone()
+ )
+ task_weights.append(task.weight)
+ return tasks, task_weights
+
- extract_task_func = get_global_func(
- "relay.backend.MetaScheduleExtractTask",
- allow_missing=False,
+def tune_relay(
+ mod: IRModule,
+ params: Dict[str, NDArray],
+ target: Union[str, Target],
+ work_dir: str,
+ max_trials_global: int,
+ *,
+ max_trials_per_task: Optional[int] = None,
+ num_trials_per_iter: int = 64,
+ builder: Builder.BuilderType = "local",
+ runner: Runner.RunnerType = "local",
+ database: Database.DatabaseType = "json",
+ cost_model: CostModel.CostModelType = "xgb",
+ measure_callbacks: MeasureCallback.CallbackListType = "default",
+ task_scheduler: TaskScheduler.TaskSchedulerType = "gradient",
+ space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+ strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+ seed: Optional[int] = None,
+) -> Database:
+ """Tune a Relay program.
+
+ Parameters
+ ----------
+ mod : Union[IRModule, tir.PrimFunc]
+ The module or function to tune
+ params : Optional[Dict[str, tvm.runtime.NDArray]]
+ The associated parameters of the program
+ target : Union[Target, str]
+ The compilation target
+ work_dir : str
+ The working directory to store the tuning records
+ max_trials_global : int
+ The maximum number of trials to run
+ max_trials_per_task : Optional[int]
+ The maximum number of trials to run for each task
+ num_trials_per_iter : int
+ The number of trials to run per iteration
+ builder : BuilderType
+ The builder to use
+ runner : RunnerType
+ The runner to use
+ database : DatabaseType
+ The database to use
+ cost_model : CostModelType
+ The cost model to use
+ measure_callbacks : CallbackListType
+ The measure callbacks to use
+ task_scheduler : TaskSchedulerType
+ The task scheduler to use
+ space : SpaceGeneratorType
+ The space generator to use
+ strategy : SearchStrategyType
+ The search strategy to use
+ seed : Optional[int]
+ The random seed
+
+ Returns
+ -------
+ database : Database
+ The database that contains the tuning records
+ """
+ tasks, task_weights = extracted_tasks_to_tune_contexts(
+ extracted_tasks=extract_tasks(mod, target, params),
Review Comment:
let me try to add an `executor` parameter to this interface
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190272
##########
python/tvm/meta_schedule/logging.py:
##########
@@ -0,0 +1,262 @@
+# 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.
+"""Logging interface in MetaSchedule"""
+import logging
+import logging.config
+import os
+import os.path as osp
+from logging import Logger
+from typing import Any, Callable, Dict, List, Optional
+
+
+def get_logger(name: str) -> Logger:
+ """Create or get a logger by its name. This is essentially a wrapper of python's native logger.
+
+ Parameters
+ ----------
+ name : str
+ The name of the logger.
+
+ Returns
+ -------
+ logger : Logger
+ The logger instance.
+ """
+ return logging.getLogger(name)
+
+
+def make_logging_func(logger: Logger) -> Optional[Callable[[int, str], None]]:
+ """Get the logging function.
+
+ Parameters
+ ----------
+ logger : Logger
+ The logger instance.
+ Returns
+ -------
+ result : Optional[Callable]
+ The function to do the specified level of logging.
+ """
+ if logger is None:
+ return None
+
+ level2log = {
+ logging.DEBUG: logger.debug,
+ logging.INFO: logger.info,
+ logging.WARNING: logger.warning,
+ logging.ERROR: logger.error,
+ # logging.FATAL not included
+ }
+
+ def logging_func(level: int, msg: str):
+ def clear_notebook_output():
+ from IPython.display import ( # type: ignore # pylint: disable=import-outside-toplevel
+ clear_output,
+ )
+
+ clear_output(wait=True)
+
+ if level < 0:
+ clear_notebook_output()
Review Comment:
> Will this be better?
Sounds good! Thanks for spotting this :-)
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1260142503
I like this change of decoupling compilation and tuning, the changes to default classes usage also make sense. Please let me know when the PR is ready for review.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989960345
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
Review Comment:
This test is broken with the error
```
> space=ms.space_generator.PostOrderApply(
f_block_filter=None,
sch_rules=None,
postprocs=[],
mutator_probs=None,
),
)
test_meta_schedule_vnni_integration.py:213:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../../python/tvm/meta_schedule/space_generator/post_order_apply.py:53: in __init__
sch_rules, postprocs, mutator_probs = _normalize_rules(sch_rules, postprocs, mutator_probs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
sch_rules = None, postprocs = [], mutator_probs = None
def _normalize_rules(
sch_rules: ScheduleRuleType,
postprocs: PostprocType,
mutator_probs: MutatorProbType,
) -> Tuple[
Optional[List["ScheduleRule"]],
Optional[List["Postproc"]],
Optional[Dict["Mutator", float]],
]:
# pylint: disable=import-outside-toplevel
from ..mutator import Mutator
from ..postproc import Postproc
from ..schedule_rule import ScheduleRule
# pylint: enable=import-outside-toplevel
> assert sch_rules is not None
E AssertionError
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r983983542
##########
python/tvm/meta_schedule/utils.py:
##########
@@ -114,7 +113,10 @@ def __init__(self, *args, **kwargs):
def __getattr__(self, name: str):
"""Bridge the attribute function."""
- return self._inst.__getattribute__(name)
+ try:
+ return self._inst.__getattribute__(name)
+ except AttributeError:
+ return super(TVMDerivedObject, self).__getattr__(name)
Review Comment:
I think it's a [known bug](https://stackoverflow.com/questions/69330379/how-to-properly-type-hint-a-class-decorator) of mypy, which haven't been fixed. I didn't see better solution than mark as ignore though.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985051216
##########
include/tvm/support/random_engine.h:
##########
@@ -46,32 +43,18 @@ namespace support {
class LinearCongruentialEngine {
public:
- /*!
- * \brief The result type is defined as uint64_t here to avoid overflow.
- * \note The type name is not in Google style because it is used in STL's distribution inferface.
Review Comment:
i dont think the doc itself is accurate
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190468
##########
include/tvm/meta_schedule/space_generator.h:
##########
@@ -171,6 +213,7 @@ class PySpaceGeneratorNode : public SpaceGeneratorNode {
FClone f_clone;
void VisitAttrs(tvm::AttrVisitor* v) {
+ SpaceGeneratorNode::VisitAttrs(v);
Review Comment:
Yep, I think so. Is there any place where I forgot to get it added?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989749128
##########
tests/python/unittest/test_meta_schedule_tune_relay.py:
##########
@@ -1,554 +0,0 @@
-# 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-docstring
-import logging
-import tempfile
-from os import path as osp
-from typing import List, Optional
-
-import numpy as np # type: ignore
-import pytest
-import tvm
-import tvm.testing
-from tvm import meta_schedule as ms
-from tvm import relay
-from tvm._ffi import register_func
-from tvm.contrib import graph_executor
-from tvm.ir import IRModule
-from tvm.meta_schedule.testing.relay_workload import get_network
-from tvm.script import tir as T
-from tvm.target.target import Target
-from tvm.tir.schedule import BlockRV, Schedule
-from tvm.tir.schedule.trace import Trace
-from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-
-logging.basicConfig(
- format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
-)
-logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
-
-# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
-# fmt: off
-
-@tvm.script.ir_module
-class tvmgen_default_fused_layout_transform:
- @T.prim_func
- def main( # type: ignore
- placeholder: T.Buffer[(1, 3, 16, 16), "float32"], # type: ignore
- T_layout_trans: T.Buffer[(1, 1, 16, 16, 3), "float32"], # type: ignore
- ) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- for i0, i1, i2, i3, i4 in T.grid(1, 1, 16, 16, 3):
- with T.block("T_layout_trans"):
- ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
- T.reads(placeholder[ax0, ax1 * 3 + ax4, ax2, ax3])
- T.writes(T_layout_trans[ax0, ax1, ax2, ax3, ax4])
- T_layout_trans[ax0, ax1, ax2, ax3, ax4] = T.if_then_else(
- ax0 < 1 and ax1 * 3 + ax4 < 3 and ax2 < 16 and ax3 < 16, # type: ignore
- placeholder[ax0, ax1 * 3 + ax4, ax2, ax3],
- T.float32(0),
- dtype="float32",
- )
-
-
-@tvm.script.ir_module
-class tvmgen_default_fused_nn_contrib_conv2d_NCHWc:
- @T.prim_func
- def main(placeholder: T.Buffer[(1, 1, 16, 16, 3), "float32"], placeholder_1: T.Buffer[(2, 1, 5, 5, 3, 4), "float32"], conv2d_NCHWc: T.Buffer[(1, 2, 16, 16, 4), "float32"]) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- data_pad = T.alloc_buffer([1, 1, 20, 20, 3], dtype="float32")
- for i0, i1, i2, i3, i4 in T.grid(1, 1, 20, 20, 3):
- with T.block("data_pad"):
- i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
- T.reads(placeholder[i0_1, i1_1, i2_1 - 2, i3_1 - 2, i4_1])
- T.writes(data_pad[i0_1, i1_1, i2_1, i3_1, i4_1])
- data_pad[i0_1, i1_1, i2_1, i3_1, i4_1] = T.if_then_else(2 <= i2_1 and i2_1 < 18 and 2 <= i3_1 and i3_1 < 18, placeholder[i0_1, i1_1, i2_1 - 2, i3_1 - 2, i4_1], T.float32(0), dtype="float32") # type: ignore # pylint: disable=R1716
- for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(1, 2, 16, 16, 4, 3, 5, 5):
- with T.block("conv2d_NCHWc"):
- n, oc_chunk, oh, ow, oc_block, ic, kh, kw = T.axis.remap("SSSSSRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
- T.reads(data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3], placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block]) # type: ignore
- T.writes(conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block])
- T.block_attr({"workload":["conv2d_NCHWc.x86", ["TENSOR", [1, 1, 16, 16, 3], "float32"], ["TENSOR", [2, 1, 5, 5, 3, 4], "float32"], [1, 1], [2, 2, 2, 2], [1, 1], "NCHW3c", "NCHW4c", "float32"]})
- with T.init():
- conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] = T.float32(0)
- conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] + data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3] * placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block] # type: ignore
-
-@tvm.script.ir_module
-class tvmgen_default_fused_layout_transform_1:
- @T.prim_func
- def main(placeholder: T.Buffer[(1, 2, 16, 16, 4), "float32"], T_layout_trans: T.Buffer[(1, 8, 16, 16), "float32"]) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- for i0, i1, i2, i3 in T.grid(1, 8, 16, 16):
- with T.block("T_layout_trans"):
- ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
- T.reads(placeholder[ax0, ax1 // 4, ax2, ax3, ax1 % 4]) # type: ignore
- T.writes(T_layout_trans[ax0, ax1, ax2, ax3])
- T_layout_trans[ax0, ax1, ax2, ax3] = T.if_then_else(ax0 < 1 and ax1 < 8 and ax2 < 16 and ax3 < 16, placeholder[ax0, ax1 // 4, ax2, ax3, ax1 % 4], T.float32(0), dtype="float32") # type: ignore
-
-# fmt: on
-# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
-
-
-@pytest.mark.skip("Integration test")
-@pytest.mark.parametrize(
- "model_name, input_shape, target, layout",
- [
- ("resnet_18", [1, 3, 224, 224], "llvm --num-cores=12", "NHWC"),
- ("resnet_18", [1, 3, 224, 224], "nvidia/geforce-rtx-3070", "NHWC"),
- ("mobilenet_v2", [1, 3, 224, 224], "llvm --num-cores=12", "NHWC"),
- ("mobilenet_v2", [1, 3, 224, 224], "nvidia/geforce-rtx-3070", "NHWC"),
- ("bert_base", [1, 64], "llvm --num-cores=12", None),
- ("bert_base", [1, 64], "nvidia/geforce-rtx-3070", None),
- ],
-)
-def test_meta_schedule_tune_relay(
- model_name: str,
- input_shape: List[int],
- target: str,
- layout: Optional[str],
-):
- dev = tvm.cpu() if str(target).startswith("llvm") else tvm.cuda()
- if model_name.startswith("bert"):
- data = tvm.nd.array(np.random.randint(0, 30521, size=input_shape), dev) # embedding size
- else:
- data = tvm.nd.array(np.random.randn(*input_shape).astype("float32"), dev)
-
- mod, params, (input_name, _, _) = get_network(
- name=model_name,
- input_shape=input_shape,
- layout=layout,
- )
-
- target = Target(target)
- with tempfile.TemporaryDirectory() as work_dir:
- with ms.Profiler() as profiler:
- rt_mod1: tvm.runtime.Module = ms.tune_relay(
- mod=mod,
- params=params,
- target=target,
- config=ms.TuneConfig(
- strategy="evolutionary",
- num_trials_per_iter=32,
- max_trials_per_task=20000,
- max_trials_global=20000,
- ),
- work_dir=work_dir,
- )
- print(profiler.table())
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def test_meta_schedule_te2primfunc_argument_order():
- @ms.derived_object
- class TestDummyDatabase(ms.database.PyDatabase):
- def __init__(self):
- super().__init__()
- self.records = []
- self.workload_reg = []
-
- def has_workload(self, mod: IRModule) -> ms.database.Workload:
- for workload in self.workload_reg:
- if tvm.ir.structural_equal(workload.mod, mod):
- return True
- # The database has already put in all correct workloads
- raise ValueError(
- "The workload searched for is not in given database!"
- + " Incorrect TIR was generated from TE subgraph."
- )
-
- def commit_tuning_record(self, record: ms.database.TuningRecord) -> None:
- self.records.append(record)
-
- def commit_workload(self, mod: IRModule) -> ms.database.Workload:
- for workload in self.workload_reg:
- if tvm.ir.structural_equal(workload.mod, mod):
- return workload
- workload = ms.database.Workload(mod)
- self.workload_reg.append(workload)
- return workload
-
- def get_top_k(
- self,
- workload: ms.database.Workload,
- top_k: int,
- ) -> List[ms.database.TuningRecord]:
- return list(
- filter(
- lambda x: x.workload == workload,
- sorted(self.records, key=lambda x: sum(x.run_secs) / len(x.run_secs)),
- )
- )[: int(top_k)]
-
- def __len__(self) -> int:
- return len(self.records)
-
- def print_results(self) -> None:
- print("\n".join([str(r) for r in self.records]))
-
- data_shape = (1, 3, 16, 16)
- weight_shape = (8, 3, 5, 5)
- data = relay.var("data", relay.TensorType(data_shape, "float32"))
- weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
- y = relay.nn.conv2d(
- data,
- weight,
- padding=(2, 2),
- kernel_size=(5, 5),
- kernel_layout="OIHW",
- out_dtype="float32",
- )
- f = relay.Function([data, weight], y)
- mod = tvm.IRModule.from_expr(f)
- mod = relay.transform.InferType()(mod)
-
- data_sample = np.random.rand(*data_shape).astype("float32")
- weight_sample = np.random.rand(*weight_shape).astype("float32")
- params = {mod["main"].params[1].name_hint: weight_sample}
-
- input_name = "data"
- dev = tvm.cpu()
- target = Target("llvm --num-cores=12")
- data = tvm.nd.array(data_sample, dev)
-
- database = TestDummyDatabase()
- database.commit_workload(tvmgen_default_fused_layout_transform)
- database.commit_workload(tvmgen_default_fused_layout_transform_1)
- database.commit_workload(tvmgen_default_fused_nn_contrib_conv2d_NCHWc)
-
- with database, tvm.transform.PassContext( # pylint: disable=not-context-manager
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- rt_mod1 = relay.build(mod, target=target, params=params)
-
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def test_meta_schedule_relay_lowering():
- data_shape = (1, 3, 16, 16)
- weight_shape = (8, 3, 5, 5)
- data = relay.var("data", relay.TensorType(data_shape, "float32"))
- weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
- y = relay.nn.conv2d(
- data,
- weight,
- padding=(2, 2),
- kernel_size=(5, 5),
- kernel_layout="OIHW",
- out_dtype="float32",
- )
- f = relay.Function([data, weight], y)
- mod = tvm.IRModule.from_expr(f)
- mod = relay.transform.InferType()(mod)
-
- data_sample = np.random.rand(*data_shape).astype("float32")
- weight_sample = np.random.rand(*weight_shape).astype("float32")
- params = {mod["main"].params[1].name_hint: weight_sample}
-
- input_name = "data"
- dev = tvm.cpu()
- target = Target("llvm --num-cores=12")
- data = tvm.nd.array(data_sample, dev)
-
- with tempfile.TemporaryDirectory() as work_dir:
- database = ms.database.JSONDatabase(
- osp.join(work_dir, "workload.json"), osp.join(work_dir, "records.json")
- )
- database.commit_tuning_record(
- ms.database.TuningRecord(
- Trace([], {}),
- database.commit_workload(tvmgen_default_fused_nn_contrib_conv2d_NCHWc),
- [0.0],
- target=target,
- args_info=[],
- )
- )
- with database, tvm.transform.PassContext(
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- rt_mod1 = relay.build(mod, target=target, params=params)
-
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def schedule_dense(dense_block, M, do_tune, sch): # pylint: disable=invalid-name
- """
- Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
- using VNNI instruction.
- """
- post_blocks = sch.get_consumers(dense_block)
-
- if len(post_blocks) > 0:
- # Fuse all intermediate post ops into the last op.
- # This is equivalent to the traverse_inline function used in TE schedules.
- while True:
- next_post_blocks = []
- for post_block in post_blocks:
- next_consumers = sch.get_consumers(post_block)
-
- if len(next_consumers) > 0:
- sch.compute_inline(post_block)
-
- next_post_blocks += next_consumers
-
- if len(next_post_blocks) == 0:
- assert len(post_blocks) == 1
- outer_block = post_blocks[0]
- a_y, a_x = sch.get_loops(outer_block)[-2:]
- break
-
- post_blocks = next_post_blocks
- else:
- a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
- outer_block = dense_block
-
- if do_tune:
- y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
- a_yo, a_yi = sch.split(a_y, factors=y_factors)
- else:
- a_yo, a_yi = sch.split(a_y, factors=[None, min(M, 64)])
-
- a_xo, a_xi = sch.split(a_x, factors=[None, 16])
- sch.reorder(a_yo, a_xo, a_yi, a_xi)
- fused = sch.fuse(a_yo, a_xo)
-
- if outer_block != dense_block:
- # Handle the case when dense is fused with post ops.
- sch.vectorize(a_xi)
- sch.compute_at(dense_block, a_yi)
-
- a_xi, a_k = sch.get_loops(dense_block)[-2:]
- a_ko, a_ki = sch.split(a_k, factors=[None, 4])
- sch.reorder(a_ko, a_xi, a_ki)
-
- # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
- # condition is violated.
- sch.parallel(fused)
- dec = sch.decompose_reduction(dense_block, a_ko)
-
- init_loop = sch.get_loops(dec)[-1]
- sch.vectorize(init_loop)
-
- sch.tensorize(a_xi, VNNI_INTRIN)
-
-
-def manual_tir_common(do_tune=False):
- M, N, K = 1024, 1024, 1024 # pylint: disable=invalid-name
- data_shape = (M, K)
- weight_shape = (N, K)
-
- data_dtype = "uint8"
- data = relay.var("data", shape=data_shape, dtype=data_dtype)
- weight = relay.var("weight", shape=weight_shape, dtype="int8")
- bias = relay.var("bias", shape=(weight_shape[0],), dtype="int32")
-
- # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
- dense = relay.nn.dense(data, weight, out_dtype="int32")
- bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
- out = relay.nn.batch_matmul(
- relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
- relay.cast(relay.expand_dims(bias_add, 0), "int8"),
- out_dtype="int32",
- )
-
- relay_mod = tvm.IRModule.from_expr(out)
-
- target = "llvm -mcpu=cascadelake -num-cores 4"
- dev = tvm.device(target, 0)
-
- data = np.random.uniform(1, 10, size=(M, K)).astype("uint8")
- weight_np = np.random.uniform(1, 10, size=weight_shape).astype("int8")
- bias_np = np.random.uniform(1, 10, size=(weight_shape[0],)).astype("int32")
-
- ref = (
- relay.create_executor("vm", mod=relay_mod, device=dev, target=target)
- .evaluate()(*[data, weight_np, bias_np])
- .numpy()
- )
-
- params = {"weight": weight_np, "bias": bias_np}
-
- if do_tune:
- extracted_tasks = ms.extract_task_from_relay(relay_mod, target, params)
- # Filter out tasks that we don't intend to schedule / tune with TIR.
- tune_tasks = list(
- filter(
- lambda task: "dense" in task.task_name,
- extracted_tasks,
- )
- )
- config = ms.TuneConfig(
- strategy="replay_trace",
- num_trials_per_iter=8,
- max_trials_per_task=8,
- max_trials_global=8,
- )
-
- with tempfile.TemporaryDirectory() as work_dir:
- # postprocs=lambda: [] is important to prevent default post processors from
- # tampering with the manual schedule.
- database = ms.tune_extracted_tasks(
- tune_tasks,
- config,
- work_dir=work_dir,
- postprocs=lambda: [],
- )
- else:
-
- def schedule_fn(sch) -> bool:
- if "dense" not in sch.mod.attrs["task_name"]:
- return False
-
- block = sch.get_block("compute")
-
- # Looks up schedule_rule annotation.
- # See the comment in test_tune_relay_manual_tir_vnni().
- schedule_rule = sch.get(block).annotations["schedule_rule"]
-
- assert "dense_vnni" in schedule_rule
-
- schedule_dense(block, M, False, sch)
-
- return True
-
- database = ms.database.ScheduleFnDatabase(schedule_fn)
-
- with database, tvm.transform.PassContext(
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- # pylint: disable=W0105
- """
- The log should say
- Warning: Cannot find workload: tvmgen_default_fused_expand_dims
- Warning: Cannot find workload: tvmgen_default_fused_cast
- Warning: Cannot find workload: tvmgen_default_fused_cast_1
- Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
-
- This means batch matmul and others are scheduled by TE, and dense (the one not warned)
- is found in the meta schedule tuning database during compilation
- """
- # pylint: enable=W0105
- lib = relay.build(relay_mod, target=target, params=params)
-
- runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
-
- runtime.set_input("data", data)
- runtime.run()
-
- out = runtime.get_output(0).numpy()
-
- np.testing.assert_equal(out, ref)
-
-
-@tvm.testing.requires_cascadelake
-def test_tune_relay_manual_tir_vnni():
Review Comment:
Sorry found the ported test in `unittest/test_meta_schedule_vnni_integration.py`. Didn't notice since the test name changed.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r990002569
##########
tests/python/contrib/test_hexagon/test_meta_schedule.py:
##########
@@ -303,23 +309,20 @@ def test_vrmpy_dense_auto_tensorize(hexagon_launcher):
if True:
with tempfile.TemporaryDirectory() as work_dir:
- config = ms.TuneConfig(
- strategy="replay_trace",
+ target = get_hexagon_target("v68")
+ database = ms.tir_integration.tune_tir(
Review Comment:
Two uses of `tune_tir` in this file have incorrect signature. I got the following errors:
```
E TypeError: tune_tir() got an unexpected keyword argument 'sch_rules'
```
```
E Check failed: (!checked_type.defined()) is false: Expected Map[meta_schedule.Mutator, FloatImm], but got Array
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985387124
##########
python/tvm/meta_schedule/task_scheduler/task_scheduler.py:
##########
@@ -101,15 +90,43 @@ def join_running_task(self, task_id: int) -> List[RunnerResult]:
"""
return _ffi_api.TaskSchedulerJoinRunningTask(self, task_id) # type: ignore # pylint: disable=no-member
- def initialize_task(self, task_id: int) -> None:
- """Initialize modules of the given task.
+ def tune(
+ self,
+ tasks: List[TuneContext],
+ task_weights: List[float],
+ max_trials_global: int,
+ max_trials_per_task: int,
+ num_trials_per_iter: int,
+ builder: Builder,
+ runner: Runner,
+ measure_callbacks: List[MeasureCallback],
+ database: Optional[Database],
+ cost_model: Optional[CostModel],
+ ) -> None:
+ """Auto-tuning."""
Review Comment:
thanks! will do!
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,82 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=False,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+) -> Tuple[IRModule, Target, Dict[str, NDArray]]:
+ from tvm.relay import Function # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+
+ return mod, target, relay_params
def extract_task_from_relay(
mod: IRModule,
- target: Target,
- params: Optional[Dict[str, NDArray]] = None,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
*,
- opt_level: int = 3,
- pass_config: Optional[Dict[str, Any]] = None,
Review Comment:
i believe as the scope of this PR, it's good that we focus on existing functionalities as tested, while get it expanded to more advanced cases could be a separate PR
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}},
Review Comment:
marked resolved as no further conversation happened in this thread - feel free to unresolve if it's not fixed :-)
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988334549
##########
python/tvm/meta_schedule/builder/builder.py:
##########
@@ -126,6 +128,22 @@ def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
"""
return _ffi_api.BuilderBuild(self, build_inputs) # type: ignore # pylint: disable=no-member
+ @staticmethod
+ def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Literal["local"] = "local",
Review Comment:
done!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988344383
##########
python/tvm/meta_schedule/database/database.py:
##########
@@ -361,6 +363,44 @@ def current() -> Optional["Database"]:
"""Get the current database under scope."""
return _ffi_api.DatabaseCurrent() # type: ignore # pylint: disable=no-member
+ @staticmethod
+ def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Union[
Review Comment:
added
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989747470
##########
tests/python/unittest/test_meta_schedule_tune_relay.py:
##########
@@ -1,554 +0,0 @@
-# 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-docstring
-import logging
-import tempfile
-from os import path as osp
-from typing import List, Optional
-
-import numpy as np # type: ignore
-import pytest
-import tvm
-import tvm.testing
-from tvm import meta_schedule as ms
-from tvm import relay
-from tvm._ffi import register_func
-from tvm.contrib import graph_executor
-from tvm.ir import IRModule
-from tvm.meta_schedule.testing.relay_workload import get_network
-from tvm.script import tir as T
-from tvm.target.target import Target
-from tvm.tir.schedule import BlockRV, Schedule
-from tvm.tir.schedule.trace import Trace
-from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-
-logging.basicConfig(
- format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
-)
-logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
-
-# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
-# fmt: off
-
-@tvm.script.ir_module
-class tvmgen_default_fused_layout_transform:
- @T.prim_func
- def main( # type: ignore
- placeholder: T.Buffer[(1, 3, 16, 16), "float32"], # type: ignore
- T_layout_trans: T.Buffer[(1, 1, 16, 16, 3), "float32"], # type: ignore
- ) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- for i0, i1, i2, i3, i4 in T.grid(1, 1, 16, 16, 3):
- with T.block("T_layout_trans"):
- ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
- T.reads(placeholder[ax0, ax1 * 3 + ax4, ax2, ax3])
- T.writes(T_layout_trans[ax0, ax1, ax2, ax3, ax4])
- T_layout_trans[ax0, ax1, ax2, ax3, ax4] = T.if_then_else(
- ax0 < 1 and ax1 * 3 + ax4 < 3 and ax2 < 16 and ax3 < 16, # type: ignore
- placeholder[ax0, ax1 * 3 + ax4, ax2, ax3],
- T.float32(0),
- dtype="float32",
- )
-
-
-@tvm.script.ir_module
-class tvmgen_default_fused_nn_contrib_conv2d_NCHWc:
- @T.prim_func
- def main(placeholder: T.Buffer[(1, 1, 16, 16, 3), "float32"], placeholder_1: T.Buffer[(2, 1, 5, 5, 3, 4), "float32"], conv2d_NCHWc: T.Buffer[(1, 2, 16, 16, 4), "float32"]) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- data_pad = T.alloc_buffer([1, 1, 20, 20, 3], dtype="float32")
- for i0, i1, i2, i3, i4 in T.grid(1, 1, 20, 20, 3):
- with T.block("data_pad"):
- i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
- T.reads(placeholder[i0_1, i1_1, i2_1 - 2, i3_1 - 2, i4_1])
- T.writes(data_pad[i0_1, i1_1, i2_1, i3_1, i4_1])
- data_pad[i0_1, i1_1, i2_1, i3_1, i4_1] = T.if_then_else(2 <= i2_1 and i2_1 < 18 and 2 <= i3_1 and i3_1 < 18, placeholder[i0_1, i1_1, i2_1 - 2, i3_1 - 2, i4_1], T.float32(0), dtype="float32") # type: ignore # pylint: disable=R1716
- for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(1, 2, 16, 16, 4, 3, 5, 5):
- with T.block("conv2d_NCHWc"):
- n, oc_chunk, oh, ow, oc_block, ic, kh, kw = T.axis.remap("SSSSSRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
- T.reads(data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3], placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block]) # type: ignore
- T.writes(conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block])
- T.block_attr({"workload":["conv2d_NCHWc.x86", ["TENSOR", [1, 1, 16, 16, 3], "float32"], ["TENSOR", [2, 1, 5, 5, 3, 4], "float32"], [1, 1], [2, 2, 2, 2], [1, 1], "NCHW3c", "NCHW4c", "float32"]})
- with T.init():
- conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] = T.float32(0)
- conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] + data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3] * placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block] # type: ignore
-
-@tvm.script.ir_module
-class tvmgen_default_fused_layout_transform_1:
- @T.prim_func
- def main(placeholder: T.Buffer[(1, 2, 16, 16, 4), "float32"], T_layout_trans: T.Buffer[(1, 8, 16, 16), "float32"]) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- for i0, i1, i2, i3 in T.grid(1, 8, 16, 16):
- with T.block("T_layout_trans"):
- ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
- T.reads(placeholder[ax0, ax1 // 4, ax2, ax3, ax1 % 4]) # type: ignore
- T.writes(T_layout_trans[ax0, ax1, ax2, ax3])
- T_layout_trans[ax0, ax1, ax2, ax3] = T.if_then_else(ax0 < 1 and ax1 < 8 and ax2 < 16 and ax3 < 16, placeholder[ax0, ax1 // 4, ax2, ax3, ax1 % 4], T.float32(0), dtype="float32") # type: ignore
-
-# fmt: on
-# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
-
-
-@pytest.mark.skip("Integration test")
-@pytest.mark.parametrize(
- "model_name, input_shape, target, layout",
- [
- ("resnet_18", [1, 3, 224, 224], "llvm --num-cores=12", "NHWC"),
- ("resnet_18", [1, 3, 224, 224], "nvidia/geforce-rtx-3070", "NHWC"),
- ("mobilenet_v2", [1, 3, 224, 224], "llvm --num-cores=12", "NHWC"),
- ("mobilenet_v2", [1, 3, 224, 224], "nvidia/geforce-rtx-3070", "NHWC"),
- ("bert_base", [1, 64], "llvm --num-cores=12", None),
- ("bert_base", [1, 64], "nvidia/geforce-rtx-3070", None),
- ],
-)
-def test_meta_schedule_tune_relay(
- model_name: str,
- input_shape: List[int],
- target: str,
- layout: Optional[str],
-):
- dev = tvm.cpu() if str(target).startswith("llvm") else tvm.cuda()
- if model_name.startswith("bert"):
- data = tvm.nd.array(np.random.randint(0, 30521, size=input_shape), dev) # embedding size
- else:
- data = tvm.nd.array(np.random.randn(*input_shape).astype("float32"), dev)
-
- mod, params, (input_name, _, _) = get_network(
- name=model_name,
- input_shape=input_shape,
- layout=layout,
- )
-
- target = Target(target)
- with tempfile.TemporaryDirectory() as work_dir:
- with ms.Profiler() as profiler:
- rt_mod1: tvm.runtime.Module = ms.tune_relay(
- mod=mod,
- params=params,
- target=target,
- config=ms.TuneConfig(
- strategy="evolutionary",
- num_trials_per_iter=32,
- max_trials_per_task=20000,
- max_trials_global=20000,
- ),
- work_dir=work_dir,
- )
- print(profiler.table())
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def test_meta_schedule_te2primfunc_argument_order():
- @ms.derived_object
- class TestDummyDatabase(ms.database.PyDatabase):
- def __init__(self):
- super().__init__()
- self.records = []
- self.workload_reg = []
-
- def has_workload(self, mod: IRModule) -> ms.database.Workload:
- for workload in self.workload_reg:
- if tvm.ir.structural_equal(workload.mod, mod):
- return True
- # The database has already put in all correct workloads
- raise ValueError(
- "The workload searched for is not in given database!"
- + " Incorrect TIR was generated from TE subgraph."
- )
-
- def commit_tuning_record(self, record: ms.database.TuningRecord) -> None:
- self.records.append(record)
-
- def commit_workload(self, mod: IRModule) -> ms.database.Workload:
- for workload in self.workload_reg:
- if tvm.ir.structural_equal(workload.mod, mod):
- return workload
- workload = ms.database.Workload(mod)
- self.workload_reg.append(workload)
- return workload
-
- def get_top_k(
- self,
- workload: ms.database.Workload,
- top_k: int,
- ) -> List[ms.database.TuningRecord]:
- return list(
- filter(
- lambda x: x.workload == workload,
- sorted(self.records, key=lambda x: sum(x.run_secs) / len(x.run_secs)),
- )
- )[: int(top_k)]
-
- def __len__(self) -> int:
- return len(self.records)
-
- def print_results(self) -> None:
- print("\n".join([str(r) for r in self.records]))
-
- data_shape = (1, 3, 16, 16)
- weight_shape = (8, 3, 5, 5)
- data = relay.var("data", relay.TensorType(data_shape, "float32"))
- weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
- y = relay.nn.conv2d(
- data,
- weight,
- padding=(2, 2),
- kernel_size=(5, 5),
- kernel_layout="OIHW",
- out_dtype="float32",
- )
- f = relay.Function([data, weight], y)
- mod = tvm.IRModule.from_expr(f)
- mod = relay.transform.InferType()(mod)
-
- data_sample = np.random.rand(*data_shape).astype("float32")
- weight_sample = np.random.rand(*weight_shape).astype("float32")
- params = {mod["main"].params[1].name_hint: weight_sample}
-
- input_name = "data"
- dev = tvm.cpu()
- target = Target("llvm --num-cores=12")
- data = tvm.nd.array(data_sample, dev)
-
- database = TestDummyDatabase()
- database.commit_workload(tvmgen_default_fused_layout_transform)
- database.commit_workload(tvmgen_default_fused_layout_transform_1)
- database.commit_workload(tvmgen_default_fused_nn_contrib_conv2d_NCHWc)
-
- with database, tvm.transform.PassContext( # pylint: disable=not-context-manager
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- rt_mod1 = relay.build(mod, target=target, params=params)
-
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def test_meta_schedule_relay_lowering():
- data_shape = (1, 3, 16, 16)
- weight_shape = (8, 3, 5, 5)
- data = relay.var("data", relay.TensorType(data_shape, "float32"))
- weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
- y = relay.nn.conv2d(
- data,
- weight,
- padding=(2, 2),
- kernel_size=(5, 5),
- kernel_layout="OIHW",
- out_dtype="float32",
- )
- f = relay.Function([data, weight], y)
- mod = tvm.IRModule.from_expr(f)
- mod = relay.transform.InferType()(mod)
-
- data_sample = np.random.rand(*data_shape).astype("float32")
- weight_sample = np.random.rand(*weight_shape).astype("float32")
- params = {mod["main"].params[1].name_hint: weight_sample}
-
- input_name = "data"
- dev = tvm.cpu()
- target = Target("llvm --num-cores=12")
- data = tvm.nd.array(data_sample, dev)
-
- with tempfile.TemporaryDirectory() as work_dir:
- database = ms.database.JSONDatabase(
- osp.join(work_dir, "workload.json"), osp.join(work_dir, "records.json")
- )
- database.commit_tuning_record(
- ms.database.TuningRecord(
- Trace([], {}),
- database.commit_workload(tvmgen_default_fused_nn_contrib_conv2d_NCHWc),
- [0.0],
- target=target,
- args_info=[],
- )
- )
- with database, tvm.transform.PassContext(
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- rt_mod1 = relay.build(mod, target=target, params=params)
-
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def schedule_dense(dense_block, M, do_tune, sch): # pylint: disable=invalid-name
- """
- Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
- using VNNI instruction.
- """
- post_blocks = sch.get_consumers(dense_block)
-
- if len(post_blocks) > 0:
- # Fuse all intermediate post ops into the last op.
- # This is equivalent to the traverse_inline function used in TE schedules.
- while True:
- next_post_blocks = []
- for post_block in post_blocks:
- next_consumers = sch.get_consumers(post_block)
-
- if len(next_consumers) > 0:
- sch.compute_inline(post_block)
-
- next_post_blocks += next_consumers
-
- if len(next_post_blocks) == 0:
- assert len(post_blocks) == 1
- outer_block = post_blocks[0]
- a_y, a_x = sch.get_loops(outer_block)[-2:]
- break
-
- post_blocks = next_post_blocks
- else:
- a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
- outer_block = dense_block
-
- if do_tune:
- y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
- a_yo, a_yi = sch.split(a_y, factors=y_factors)
- else:
- a_yo, a_yi = sch.split(a_y, factors=[None, min(M, 64)])
-
- a_xo, a_xi = sch.split(a_x, factors=[None, 16])
- sch.reorder(a_yo, a_xo, a_yi, a_xi)
- fused = sch.fuse(a_yo, a_xo)
-
- if outer_block != dense_block:
- # Handle the case when dense is fused with post ops.
- sch.vectorize(a_xi)
- sch.compute_at(dense_block, a_yi)
-
- a_xi, a_k = sch.get_loops(dense_block)[-2:]
- a_ko, a_ki = sch.split(a_k, factors=[None, 4])
- sch.reorder(a_ko, a_xi, a_ki)
-
- # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
- # condition is violated.
- sch.parallel(fused)
- dec = sch.decompose_reduction(dense_block, a_ko)
-
- init_loop = sch.get_loops(dec)[-1]
- sch.vectorize(init_loop)
-
- sch.tensorize(a_xi, VNNI_INTRIN)
-
-
-def manual_tir_common(do_tune=False):
- M, N, K = 1024, 1024, 1024 # pylint: disable=invalid-name
- data_shape = (M, K)
- weight_shape = (N, K)
-
- data_dtype = "uint8"
- data = relay.var("data", shape=data_shape, dtype=data_dtype)
- weight = relay.var("weight", shape=weight_shape, dtype="int8")
- bias = relay.var("bias", shape=(weight_shape[0],), dtype="int32")
-
- # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
- dense = relay.nn.dense(data, weight, out_dtype="int32")
- bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
- out = relay.nn.batch_matmul(
- relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
- relay.cast(relay.expand_dims(bias_add, 0), "int8"),
- out_dtype="int32",
- )
-
- relay_mod = tvm.IRModule.from_expr(out)
-
- target = "llvm -mcpu=cascadelake -num-cores 4"
- dev = tvm.device(target, 0)
-
- data = np.random.uniform(1, 10, size=(M, K)).astype("uint8")
- weight_np = np.random.uniform(1, 10, size=weight_shape).astype("int8")
- bias_np = np.random.uniform(1, 10, size=(weight_shape[0],)).astype("int32")
-
- ref = (
- relay.create_executor("vm", mod=relay_mod, device=dev, target=target)
- .evaluate()(*[data, weight_np, bias_np])
- .numpy()
- )
-
- params = {"weight": weight_np, "bias": bias_np}
-
- if do_tune:
- extracted_tasks = ms.extract_task_from_relay(relay_mod, target, params)
- # Filter out tasks that we don't intend to schedule / tune with TIR.
- tune_tasks = list(
- filter(
- lambda task: "dense" in task.task_name,
- extracted_tasks,
- )
- )
- config = ms.TuneConfig(
- strategy="replay_trace",
- num_trials_per_iter=8,
- max_trials_per_task=8,
- max_trials_global=8,
- )
-
- with tempfile.TemporaryDirectory() as work_dir:
- # postprocs=lambda: [] is important to prevent default post processors from
- # tampering with the manual schedule.
- database = ms.tune_extracted_tasks(
- tune_tasks,
- config,
- work_dir=work_dir,
- postprocs=lambda: [],
- )
- else:
-
- def schedule_fn(sch) -> bool:
- if "dense" not in sch.mod.attrs["task_name"]:
- return False
-
- block = sch.get_block("compute")
-
- # Looks up schedule_rule annotation.
- # See the comment in test_tune_relay_manual_tir_vnni().
- schedule_rule = sch.get(block).annotations["schedule_rule"]
-
- assert "dense_vnni" in schedule_rule
-
- schedule_dense(block, M, False, sch)
-
- return True
-
- database = ms.database.ScheduleFnDatabase(schedule_fn)
-
- with database, tvm.transform.PassContext(
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- # pylint: disable=W0105
- """
- The log should say
- Warning: Cannot find workload: tvmgen_default_fused_expand_dims
- Warning: Cannot find workload: tvmgen_default_fused_cast
- Warning: Cannot find workload: tvmgen_default_fused_cast_1
- Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
-
- This means batch matmul and others are scheduled by TE, and dense (the one not warned)
- is found in the meta schedule tuning database during compilation
- """
- # pylint: enable=W0105
- lib = relay.build(relay_mod, target=target, params=params)
-
- runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
-
- runtime.set_input("data", data)
- runtime.run()
-
- out = runtime.get_output(0).numpy()
-
- np.testing.assert_equal(out, ref)
-
-
-@tvm.testing.requires_cascadelake
-def test_tune_relay_manual_tir_vnni():
Review Comment:
I migrated them to integration test folder to be more consistent
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988620495
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,101 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+if TYPE_CHECKING:
+ from tvm import relay
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=True,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+ pass_config: Mapping[str, Any],
+ executor: Optional["relay.backend.Executor"],
+) -> Tuple[IRModule, Target, Dict[str, NDArray], Dict[str, Any], "relay.backend.Executor"]:
+ from tvm import relay # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, relay.Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+ if executor is None:
+ executor = relay.backend.Executor("graph")
+ pass_config = dict(pass_config)
+ pass_config.setdefault(
+ "relay.FuseOps.link_params",
+ executor.attrs.get("link_params", False),
+ )
Review Comment:
Instead of modifying `pass_config`, we can use `mod.with_attr("executor", executor)` similarly to https://github.com/apache/tvm/blob/d4bf9ecf5524d265916ac7b860b0027f5eee5c49/src/relay/backend/build_module.cc#L414-L415. This lets us align `link-params` attribute lookup in task extraction and `relay.build` (instead of relying on the ad hoc config `relay.FuseOps.link_params` for the former.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985899276
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,82 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=False,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+) -> Tuple[IRModule, Target, Dict[str, NDArray]]:
+ from tvm.relay import Function # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+
+ return mod, target, relay_params
def extract_task_from_relay(
mod: IRModule,
- target: Target,
- params: Optional[Dict[str, NDArray]] = None,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
*,
- opt_level: int = 3,
- pass_config: Optional[Dict[str, Any]] = None,
Review Comment:
i changed my mind - will add those APIs so that the latest commit (https://github.com/apache/tvm/pull/12854/) won't be broken
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,82 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=False,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+) -> Tuple[IRModule, Target, Dict[str, NDArray]]:
+ from tvm.relay import Function # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+
+ return mod, target, relay_params
def extract_task_from_relay(
mod: IRModule,
- target: Target,
- params: Optional[Dict[str, NDArray]] = None,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
*,
- opt_level: int = 3,
- pass_config: Optional[Dict[str, Any]] = None,
Review Comment:
changed my mind - will add those APIs so that the latest commit (fa17da22c73fb9e95c27e4c28130835b628caf6b) won't be broken
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988367063
##########
python/tvm/meta_schedule/schedule_rule/schedule_rule.py:
##########
@@ -76,6 +81,35 @@ def clone(self) -> "ScheduleRule":
"""
return _ffi_api.ScheduleRuleClone(self) # type: ignore # pylint: disable=no-member
+ @staticmethod
+ def create(kind: Literal["llvm", "cuda", "cuda-tensorcore"]) -> List["ScheduleRule"]:
Review Comment:
added
##########
python/tvm/meta_schedule/mutator/mutator.py:
##########
@@ -68,6 +73,35 @@ def clone(self) -> "Mutator":
"""
return _ffi_api.MutatorClone(self) # type: ignore # pylint: disable=no-member
+ @staticmethod
+ def create(kind: Literal["llvm", "cuda", "cuda-tensorcore"]) -> Dict["Mutator", float]:
Review Comment:
added
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
vinx13 commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r984973110
##########
tests/python/unittest/test_meta_schedule_space_generator.py:
##########
@@ -104,4 +108,5 @@ class TestPySpaceGenerator(PySpaceGenerator):
if __name__ == "__main__":
- tvm.testing.main()
+ test_meta_schedule_design_space_generator_NIE()
+ # tvm.testing.main()
Review Comment:
keep this
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
Review Comment:
```suggestion
/*max_innermost_factor=*/Integer(4),
```
##########
src/meta_schedule/task_scheduler/task_scheduler.cc:
##########
@@ -21,83 +21,225 @@
namespace tvm {
namespace meta_schedule {
-void TaskSchedulerNode::InitializeTask(int task_id) {
+TaskRecord::TaskRecord(TuneContext ctx, double task_weight) {
+ ObjectPtr<TaskRecordNode> n = runtime::make_object<TaskRecordNode>();
+ n->ctx = ctx;
+ n->task_weight = task_weight;
+ n->flop = 1.0;
auto _ = Profiler::TimedScope("InitializeTask");
- TuneContext task = this->tasks[task_id];
- TVM_PY_LOG(INFO, this->logging_func)
- << "Initializing Task #" << task_id << ": " << task->task_name;
- TVM_PY_LOG(INFO, task->logging_func)
- << "Initializing Task #" << task_id << ": " << task->task_name;
- CHECK(task->mod.defined()) << "ValueError: Require `context.mod`, but it is not defined";
- CHECK(task->space_generator.defined())
+ CHECK(ctx->mod.defined()) << "ValueError: Require `context.mod`, but it is not defined";
+ CHECK(ctx->space_generator.defined())
<< "ValueError: Require `context.space_generator`, but it is not defined";
- CHECK(task->search_strategy.defined())
+ CHECK(ctx->search_strategy.defined())
<< "ValueError: Require `context.search_strategy`, but it is not defined";
- TVM_PY_LOG(INFO, task->logging_func) << "\n" << tir::AsTVMScript(task->mod);
- task->Initialize();
- Array<tir::Schedule> design_spaces =
- task->space_generator.value()->GenerateDesignSpace(task->mod.value());
- TVM_PY_LOG(INFO, task->logging_func)
- << "Total " << design_spaces.size() << " design space(s) generated";
- for (int i = 0, n = design_spaces.size(); i < n; ++i) {
- tir::Schedule sch = design_spaces[i];
- tir::Trace trace = sch->trace().value();
- trace = trace->Simplified(true);
- TVM_PY_LOG(INFO, task->logging_func) << "Design space #" << i << ":\n"
- << tir::AsTVMScript(sch->mod()) << "\n"
- << Concat(trace->AsPython(false), "\n");
+ TVM_PY_LOG(INFO, ctx->logger) << "\n" << tir::AsTVMScript(ctx->mod);
+ ctx->Initialize();
+ n->flop = std::max(1.0, tir::EstimateTIRFlops(ctx->mod.value()));
+ this->data_ = std::move(n);
+}
+
+void SendToBuilder(TaskRecordNode* self, const Builder& builder) {
+ auto _ = Profiler::TimedScope("SendToBuilder");
+ Array<MeasureCandidate> candidates = self->measure_candidates.value();
+ Target target = self->ctx->target.value();
+ Array<BuilderInput> inputs;
+ inputs.reserve(candidates.size());
+ for (const MeasureCandidate& candidate : candidates) {
+ inputs.push_back(BuilderInput(candidate->sch->mod(), target));
}
- task->search_strategy.value()->PreTuning(design_spaces, database, cost_model);
+ self->builder_results = builder->Build(inputs);
}
-void TaskSchedulerNode::Tune() {
- int n_tasks = this->tasks.size();
- for (int task_id = 0; task_id < n_tasks; ++task_id) {
- InitializeTask(task_id);
+void SendToRunner(TaskRecordNode* self, const Runner& runner) {
+ auto _ = Profiler::TimedScope("SendToRunner");
+ Array<MeasureCandidate> candidates = self->measure_candidates.value();
+ Array<BuilderResult> builder_results = self->builder_results.value();
+ Target target = self->ctx->target.value();
+ ICHECK_EQ(candidates.size(), builder_results.size());
+ int n = candidates.size();
+ int n_build_errors = 0;
+ Array<RunnerInput> inputs;
+ inputs.reserve(n);
+ for (int i = 0; i < n; ++i) {
+ const MeasureCandidate& candidate = candidates[i];
+ const BuilderResult& builder_result = builder_results[i];
+ if (builder_result->error_msg.defined()) {
+ ++n_build_errors;
+ continue;
+ }
+ inputs.push_back(RunnerInput(/*artifact_path=*/builder_result->artifact_path.value(),
+ /*device_type=*/target->kind->name,
+ /*args_info=*/candidate->args_info));
+ }
+ Array<RunnerFuture> futures = runner->Run(inputs);
+ if (n_build_errors == 0) {
+ self->runner_futures = futures;
+ return;
+ }
+ Array<RunnerFuture> results;
+ results.reserve(n);
+ for (int i = 0, j = 0; i < n; ++i) {
+ const BuilderResult& builder_result = builder_results[i];
+ if (builder_result->error_msg.defined()) {
+ results.push_back(RunnerFuture(
+ /*f_done=*/[]() -> bool { return true; },
+ /*f_result=*/
+ [msg = builder_result->error_msg]() -> RunnerResult {
+ return RunnerResult(NullOpt, msg);
+ }));
+ } else {
+ results.push_back(futures[j++]);
+ }
+ }
+ self->runner_futures = results;
+}
+
+void TaskCleanUp(TaskRecordNode* self, int task_id, const Array<RunnerResult>& results) {
+ ICHECK_EQ(self->builder_results.value().size(), results.size());
+ ICHECK_EQ(self->runner_futures.value().size(), results.size());
+ int n = results.size();
+ std::string name = self->ctx->task_name.value();
+ const PackedFunc& logger = self->ctx->logger;
+ for (int i = 0; i < n; ++i) {
+ const BuilderResult& builder_result = self->builder_results.value()[i];
+ const MeasureCandidate& candidate = self->measure_candidates.value()[i];
+ const RunnerResult& runner_result = results[i];
+ Optional<String> error_msg = NullOpt;
+ int trials = self->latency_ms.size() + 1;
+ double run_ms = 1e9;
+ if ((error_msg = builder_result->error_msg)) {
+ ++self->build_error_count;
+ } else if ((error_msg = runner_result->error_msg)) {
+ ++self->run_error_count;
+ } else {
+ run_ms = GetRunMsMedian(runner_result);
+ }
+ self->latency_ms.push_back(run_ms);
+ if (error_msg) {
Review Comment:
I'm wondering how frequent these errors happen during tuning. Should we print all errors or just errors above count threshold?
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,82 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=False,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+) -> Tuple[IRModule, Target, Dict[str, NDArray]]:
+ from tvm.relay import Function # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+
+ return mod, target, relay_params
def extract_task_from_relay(
mod: IRModule,
- target: Target,
- params: Optional[Dict[str, NDArray]] = None,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
*,
- opt_level: int = 3,
- pass_config: Optional[Dict[str, Any]] = None,
Review Comment:
do we have alternative? a possible use case is to tune and compile relay under a different pass config or opt level
##########
python/tvm/meta_schedule/task_scheduler/task_scheduler.py:
##########
@@ -101,15 +90,43 @@ def join_running_task(self, task_id: int) -> List[RunnerResult]:
"""
return _ffi_api.TaskSchedulerJoinRunningTask(self, task_id) # type: ignore # pylint: disable=no-member
- def initialize_task(self, task_id: int) -> None:
- """Initialize modules of the given task.
+ def tune(
+ self,
+ tasks: List[TuneContext],
+ task_weights: List[float],
+ max_trials_global: int,
+ max_trials_per_task: int,
+ num_trials_per_iter: int,
+ builder: Builder,
+ runner: Runner,
+ measure_callbacks: List[MeasureCallback],
+ database: Optional[Database],
+ cost_model: Optional[CostModel],
+ ) -> None:
+ """Auto-tuning."""
Review Comment:
document params
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190584
##########
include/tvm/meta_schedule/task_scheduler.h:
##########
@@ -73,66 +127,71 @@ namespace meta_schedule {
*/
class TaskSchedulerNode : public runtime::Object {
public:
- /*! \brief The tasks to be tuned */
- Array<TuneContext> tasks;
- /*! \brief The builder of the scheduler. */
- Builder builder{nullptr};
- /*! \brief The runner of the scheduler. */
- Runner runner{nullptr};
- /*! \brief The database of the scheduler. */
- Optional<Database> database;
- /*! \brief The cost model of the scheduler. */
- Optional<CostModel> cost_model;
+ /*! \brief The tuning task's logging function. */
+ PackedFunc logger;
+ /*! \brief Records for each task */
+ Array<TaskRecord> tasks_;
/*! \brief The list of measure callbacks of the scheduler. */
- Array<MeasureCallback> measure_callbacks;
- /*! \brief The maximum number of trials allowed. */
- int max_trials;
- /*! \brief The number of trials already conducted. */
- int num_trials_already;
- /*! \brief The tuning task's logging function. t*/
- PackedFunc logging_func;
+ Array<MeasureCallback> measure_callbacks_;
+ /*! \brief The database used in tuning */
+ Optional<Database> database_;
+ /*! \brief The cost model used in tuning */
+ Optional<CostModel> cost_model_;
+ /*! \brief The number of remaining tasks to be tuned. */
+ int remaining_tasks_;
/*! \brief The default destructor. */
virtual ~TaskSchedulerNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) {
- v->Visit("tasks", &tasks);
- v->Visit("builder", &builder);
- v->Visit("runner", &runner);
- v->Visit("database", &database);
- v->Visit("cost_model", &cost_model);
- v->Visit("measure_callbacks", &measure_callbacks);
- v->Visit("max_trials", &max_trials);
- v->Visit("num_trials_already", &num_trials_already);
- // `logging_func` is not visited
+ // `logger` is not visited
+ v->Visit("tasks_", &tasks_);
+ v->Visit("measure_callbacks_", &measure_callbacks_);
+ v->Visit("database_", &database_);
+ v->Visit("cost_model_", &cost_model_);
+ v->Visit("remaining_tasks_", &remaining_tasks_);
}
- /*! \brief Auto-tuning. */
- virtual void Tune();
-
- /*!
- * \brief Initialize modules of the given task.
- * \param task_id The task id to be initialized.
- */
- virtual void InitializeTask(int task_id);
-
/*!
- * \brief Touch the task and update its status
- * \param task_id The task id to be checked.
+ * \brief Fetch the next task id.
+ * \return The next task id.
*/
- virtual void TouchTask(int task_id);
-
+ virtual int NextTaskId() = 0;
/*!
* \brief Wait until the task is finished.
* \param task_id The task id to be joined.
+ * \return The results from the runner.
*/
virtual Array<RunnerResult> JoinRunningTask(int task_id);
-
/*!
- * \brief Fetch the next task id.
- * \return The next task id.
+ * \brief Jointly tune a given list of tasks.
+ * \param tasks The tasks to be tuned
+ * \param task_weights The weight of each task
+ * \param max_trials_global The maximum number of trials to be performed globally
+ * \param max_trials_per_task The maximum number of trials to be performed for each task
+ * \param num_trials_per_iter The number of trials to be performed in each iteration
+ * \param builder The MetaSchedule builder
+ * \param runner The MetaSchedule runner
+ * \param measure_callbacks The callbacks to be called after each measurement
+ * \param database The database used in tuning
+ * \param cost_model The cost model used in tuning
*/
- virtual int NextTaskId() = 0;
+ virtual void Tune(Array<TuneContext> tasks, Array<FloatImm> task_weights, int max_trials_global,
Review Comment:
Well, sure! Certainly it's personal preference but indeed your suggestion makes sense
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190975
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,82 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Tuple, Union
+# isort: off
Review Comment:
I feel weird about this too, but tried all possible ways and didn't find any working to disable isort
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
tqchen commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1257191748
great to see clear separation between tuning and apply. A positive step towards modular automation
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
vinx13 commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r982956907
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}},
Review Comment:
see https://github.com/apache/tvm/blob/main/python/tvm/meta_schedule/default_config.py#L374-L379
it's different from non-tensor core config
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989573744
##########
tests/python/contrib/test_hexagon/test_meta_schedule.py:
##########
@@ -358,6 +361,7 @@ def test_conv2d_relay_auto_schedule(hexagon_launcher):
kernel_layout="HWIO",
)
mod = tvm.IRModule.from_expr(conv2d + bias)
+ mod = mod.with_attr("executor", relay.backend.Executor("graph"))
Review Comment:
We need `executor = Executor("graph", {"link-params": True})` 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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1270867070
@junrushao I made one comment but otherwise Hexagon change looks good to me. It didn't occur to me before that we can do `mod = mod.with_attr(...)` from the user script to avoid threading `executor` through task extraction and `tune_relay` etc.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989746409
##########
tests/python/unittest/test_meta_schedule_tune_relay.py:
##########
@@ -1,554 +0,0 @@
-# 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-docstring
-import logging
-import tempfile
-from os import path as osp
-from typing import List, Optional
-
-import numpy as np # type: ignore
-import pytest
-import tvm
-import tvm.testing
-from tvm import meta_schedule as ms
-from tvm import relay
-from tvm._ffi import register_func
-from tvm.contrib import graph_executor
-from tvm.ir import IRModule
-from tvm.meta_schedule.testing.relay_workload import get_network
-from tvm.script import tir as T
-from tvm.target.target import Target
-from tvm.tir.schedule import BlockRV, Schedule
-from tvm.tir.schedule.trace import Trace
-from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-
-logging.basicConfig(
- format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
-)
-logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
-
-# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
-# fmt: off
-
-@tvm.script.ir_module
-class tvmgen_default_fused_layout_transform:
- @T.prim_func
- def main( # type: ignore
- placeholder: T.Buffer[(1, 3, 16, 16), "float32"], # type: ignore
- T_layout_trans: T.Buffer[(1, 1, 16, 16, 3), "float32"], # type: ignore
- ) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- for i0, i1, i2, i3, i4 in T.grid(1, 1, 16, 16, 3):
- with T.block("T_layout_trans"):
- ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
- T.reads(placeholder[ax0, ax1 * 3 + ax4, ax2, ax3])
- T.writes(T_layout_trans[ax0, ax1, ax2, ax3, ax4])
- T_layout_trans[ax0, ax1, ax2, ax3, ax4] = T.if_then_else(
- ax0 < 1 and ax1 * 3 + ax4 < 3 and ax2 < 16 and ax3 < 16, # type: ignore
- placeholder[ax0, ax1 * 3 + ax4, ax2, ax3],
- T.float32(0),
- dtype="float32",
- )
-
-
-@tvm.script.ir_module
-class tvmgen_default_fused_nn_contrib_conv2d_NCHWc:
- @T.prim_func
- def main(placeholder: T.Buffer[(1, 1, 16, 16, 3), "float32"], placeholder_1: T.Buffer[(2, 1, 5, 5, 3, 4), "float32"], conv2d_NCHWc: T.Buffer[(1, 2, 16, 16, 4), "float32"]) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- data_pad = T.alloc_buffer([1, 1, 20, 20, 3], dtype="float32")
- for i0, i1, i2, i3, i4 in T.grid(1, 1, 20, 20, 3):
- with T.block("data_pad"):
- i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
- T.reads(placeholder[i0_1, i1_1, i2_1 - 2, i3_1 - 2, i4_1])
- T.writes(data_pad[i0_1, i1_1, i2_1, i3_1, i4_1])
- data_pad[i0_1, i1_1, i2_1, i3_1, i4_1] = T.if_then_else(2 <= i2_1 and i2_1 < 18 and 2 <= i3_1 and i3_1 < 18, placeholder[i0_1, i1_1, i2_1 - 2, i3_1 - 2, i4_1], T.float32(0), dtype="float32") # type: ignore # pylint: disable=R1716
- for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(1, 2, 16, 16, 4, 3, 5, 5):
- with T.block("conv2d_NCHWc"):
- n, oc_chunk, oh, ow, oc_block, ic, kh, kw = T.axis.remap("SSSSSRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
- T.reads(data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3], placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block]) # type: ignore
- T.writes(conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block])
- T.block_attr({"workload":["conv2d_NCHWc.x86", ["TENSOR", [1, 1, 16, 16, 3], "float32"], ["TENSOR", [2, 1, 5, 5, 3, 4], "float32"], [1, 1], [2, 2, 2, 2], [1, 1], "NCHW3c", "NCHW4c", "float32"]})
- with T.init():
- conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] = T.float32(0)
- conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc[n, oc_chunk, oh, ow, oc_block] + data_pad[n, ic // 3, oh + kh, ow + kw, ic % 3] * placeholder_1[oc_chunk, ic // 3, kh, kw, ic % 3, oc_block] # type: ignore
-
-@tvm.script.ir_module
-class tvmgen_default_fused_layout_transform_1:
- @T.prim_func
- def main(placeholder: T.Buffer[(1, 2, 16, 16, 4), "float32"], T_layout_trans: T.Buffer[(1, 8, 16, 16), "float32"]) -> None: # type: ignore
- # function attr dict
- T.func_attr({"global_symbol": "main", "tir.noalias": True})
- # body
- # with T.block("root")
- for i0, i1, i2, i3 in T.grid(1, 8, 16, 16):
- with T.block("T_layout_trans"):
- ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
- T.reads(placeholder[ax0, ax1 // 4, ax2, ax3, ax1 % 4]) # type: ignore
- T.writes(T_layout_trans[ax0, ax1, ax2, ax3])
- T_layout_trans[ax0, ax1, ax2, ax3] = T.if_then_else(ax0 < 1 and ax1 < 8 and ax2 < 16 and ax3 < 16, placeholder[ax0, ax1 // 4, ax2, ax3, ax1 % 4], T.float32(0), dtype="float32") # type: ignore
-
-# fmt: on
-# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
-
-
-@pytest.mark.skip("Integration test")
-@pytest.mark.parametrize(
- "model_name, input_shape, target, layout",
- [
- ("resnet_18", [1, 3, 224, 224], "llvm --num-cores=12", "NHWC"),
- ("resnet_18", [1, 3, 224, 224], "nvidia/geforce-rtx-3070", "NHWC"),
- ("mobilenet_v2", [1, 3, 224, 224], "llvm --num-cores=12", "NHWC"),
- ("mobilenet_v2", [1, 3, 224, 224], "nvidia/geforce-rtx-3070", "NHWC"),
- ("bert_base", [1, 64], "llvm --num-cores=12", None),
- ("bert_base", [1, 64], "nvidia/geforce-rtx-3070", None),
- ],
-)
-def test_meta_schedule_tune_relay(
- model_name: str,
- input_shape: List[int],
- target: str,
- layout: Optional[str],
-):
- dev = tvm.cpu() if str(target).startswith("llvm") else tvm.cuda()
- if model_name.startswith("bert"):
- data = tvm.nd.array(np.random.randint(0, 30521, size=input_shape), dev) # embedding size
- else:
- data = tvm.nd.array(np.random.randn(*input_shape).astype("float32"), dev)
-
- mod, params, (input_name, _, _) = get_network(
- name=model_name,
- input_shape=input_shape,
- layout=layout,
- )
-
- target = Target(target)
- with tempfile.TemporaryDirectory() as work_dir:
- with ms.Profiler() as profiler:
- rt_mod1: tvm.runtime.Module = ms.tune_relay(
- mod=mod,
- params=params,
- target=target,
- config=ms.TuneConfig(
- strategy="evolutionary",
- num_trials_per_iter=32,
- max_trials_per_task=20000,
- max_trials_global=20000,
- ),
- work_dir=work_dir,
- )
- print(profiler.table())
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def test_meta_schedule_te2primfunc_argument_order():
- @ms.derived_object
- class TestDummyDatabase(ms.database.PyDatabase):
- def __init__(self):
- super().__init__()
- self.records = []
- self.workload_reg = []
-
- def has_workload(self, mod: IRModule) -> ms.database.Workload:
- for workload in self.workload_reg:
- if tvm.ir.structural_equal(workload.mod, mod):
- return True
- # The database has already put in all correct workloads
- raise ValueError(
- "The workload searched for is not in given database!"
- + " Incorrect TIR was generated from TE subgraph."
- )
-
- def commit_tuning_record(self, record: ms.database.TuningRecord) -> None:
- self.records.append(record)
-
- def commit_workload(self, mod: IRModule) -> ms.database.Workload:
- for workload in self.workload_reg:
- if tvm.ir.structural_equal(workload.mod, mod):
- return workload
- workload = ms.database.Workload(mod)
- self.workload_reg.append(workload)
- return workload
-
- def get_top_k(
- self,
- workload: ms.database.Workload,
- top_k: int,
- ) -> List[ms.database.TuningRecord]:
- return list(
- filter(
- lambda x: x.workload == workload,
- sorted(self.records, key=lambda x: sum(x.run_secs) / len(x.run_secs)),
- )
- )[: int(top_k)]
-
- def __len__(self) -> int:
- return len(self.records)
-
- def print_results(self) -> None:
- print("\n".join([str(r) for r in self.records]))
-
- data_shape = (1, 3, 16, 16)
- weight_shape = (8, 3, 5, 5)
- data = relay.var("data", relay.TensorType(data_shape, "float32"))
- weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
- y = relay.nn.conv2d(
- data,
- weight,
- padding=(2, 2),
- kernel_size=(5, 5),
- kernel_layout="OIHW",
- out_dtype="float32",
- )
- f = relay.Function([data, weight], y)
- mod = tvm.IRModule.from_expr(f)
- mod = relay.transform.InferType()(mod)
-
- data_sample = np.random.rand(*data_shape).astype("float32")
- weight_sample = np.random.rand(*weight_shape).astype("float32")
- params = {mod["main"].params[1].name_hint: weight_sample}
-
- input_name = "data"
- dev = tvm.cpu()
- target = Target("llvm --num-cores=12")
- data = tvm.nd.array(data_sample, dev)
-
- database = TestDummyDatabase()
- database.commit_workload(tvmgen_default_fused_layout_transform)
- database.commit_workload(tvmgen_default_fused_layout_transform_1)
- database.commit_workload(tvmgen_default_fused_nn_contrib_conv2d_NCHWc)
-
- with database, tvm.transform.PassContext( # pylint: disable=not-context-manager
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- rt_mod1 = relay.build(mod, target=target, params=params)
-
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def test_meta_schedule_relay_lowering():
- data_shape = (1, 3, 16, 16)
- weight_shape = (8, 3, 5, 5)
- data = relay.var("data", relay.TensorType(data_shape, "float32"))
- weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
- y = relay.nn.conv2d(
- data,
- weight,
- padding=(2, 2),
- kernel_size=(5, 5),
- kernel_layout="OIHW",
- out_dtype="float32",
- )
- f = relay.Function([data, weight], y)
- mod = tvm.IRModule.from_expr(f)
- mod = relay.transform.InferType()(mod)
-
- data_sample = np.random.rand(*data_shape).astype("float32")
- weight_sample = np.random.rand(*weight_shape).astype("float32")
- params = {mod["main"].params[1].name_hint: weight_sample}
-
- input_name = "data"
- dev = tvm.cpu()
- target = Target("llvm --num-cores=12")
- data = tvm.nd.array(data_sample, dev)
-
- with tempfile.TemporaryDirectory() as work_dir:
- database = ms.database.JSONDatabase(
- osp.join(work_dir, "workload.json"), osp.join(work_dir, "records.json")
- )
- database.commit_tuning_record(
- ms.database.TuningRecord(
- Trace([], {}),
- database.commit_workload(tvmgen_default_fused_nn_contrib_conv2d_NCHWc),
- [0.0],
- target=target,
- args_info=[],
- )
- )
- with database, tvm.transform.PassContext(
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- rt_mod1 = relay.build(mod, target=target, params=params)
-
- # Compile without meta-schedule for correctness check
- with tvm.transform.PassContext(opt_level=0):
- rt_mod2 = relay.build(mod, target=target, params=params)
-
- def get_output(data, lib):
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(input_name, data)
- module.run()
- return module.get_output(0).numpy()
-
- # Check correctness
- actual_output = get_output(data, rt_mod1)
- expected_output = get_output(data, rt_mod2)
- assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
-
-
-def schedule_dense(dense_block, M, do_tune, sch): # pylint: disable=invalid-name
- """
- Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
- using VNNI instruction.
- """
- post_blocks = sch.get_consumers(dense_block)
-
- if len(post_blocks) > 0:
- # Fuse all intermediate post ops into the last op.
- # This is equivalent to the traverse_inline function used in TE schedules.
- while True:
- next_post_blocks = []
- for post_block in post_blocks:
- next_consumers = sch.get_consumers(post_block)
-
- if len(next_consumers) > 0:
- sch.compute_inline(post_block)
-
- next_post_blocks += next_consumers
-
- if len(next_post_blocks) == 0:
- assert len(post_blocks) == 1
- outer_block = post_blocks[0]
- a_y, a_x = sch.get_loops(outer_block)[-2:]
- break
-
- post_blocks = next_post_blocks
- else:
- a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
- outer_block = dense_block
-
- if do_tune:
- y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
- a_yo, a_yi = sch.split(a_y, factors=y_factors)
- else:
- a_yo, a_yi = sch.split(a_y, factors=[None, min(M, 64)])
-
- a_xo, a_xi = sch.split(a_x, factors=[None, 16])
- sch.reorder(a_yo, a_xo, a_yi, a_xi)
- fused = sch.fuse(a_yo, a_xo)
-
- if outer_block != dense_block:
- # Handle the case when dense is fused with post ops.
- sch.vectorize(a_xi)
- sch.compute_at(dense_block, a_yi)
-
- a_xi, a_k = sch.get_loops(dense_block)[-2:]
- a_ko, a_ki = sch.split(a_k, factors=[None, 4])
- sch.reorder(a_ko, a_xi, a_ki)
-
- # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
- # condition is violated.
- sch.parallel(fused)
- dec = sch.decompose_reduction(dense_block, a_ko)
-
- init_loop = sch.get_loops(dec)[-1]
- sch.vectorize(init_loop)
-
- sch.tensorize(a_xi, VNNI_INTRIN)
-
-
-def manual_tir_common(do_tune=False):
- M, N, K = 1024, 1024, 1024 # pylint: disable=invalid-name
- data_shape = (M, K)
- weight_shape = (N, K)
-
- data_dtype = "uint8"
- data = relay.var("data", shape=data_shape, dtype=data_dtype)
- weight = relay.var("weight", shape=weight_shape, dtype="int8")
- bias = relay.var("bias", shape=(weight_shape[0],), dtype="int32")
-
- # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
- dense = relay.nn.dense(data, weight, out_dtype="int32")
- bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
- out = relay.nn.batch_matmul(
- relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
- relay.cast(relay.expand_dims(bias_add, 0), "int8"),
- out_dtype="int32",
- )
-
- relay_mod = tvm.IRModule.from_expr(out)
-
- target = "llvm -mcpu=cascadelake -num-cores 4"
- dev = tvm.device(target, 0)
-
- data = np.random.uniform(1, 10, size=(M, K)).astype("uint8")
- weight_np = np.random.uniform(1, 10, size=weight_shape).astype("int8")
- bias_np = np.random.uniform(1, 10, size=(weight_shape[0],)).astype("int32")
-
- ref = (
- relay.create_executor("vm", mod=relay_mod, device=dev, target=target)
- .evaluate()(*[data, weight_np, bias_np])
- .numpy()
- )
-
- params = {"weight": weight_np, "bias": bias_np}
-
- if do_tune:
- extracted_tasks = ms.extract_task_from_relay(relay_mod, target, params)
- # Filter out tasks that we don't intend to schedule / tune with TIR.
- tune_tasks = list(
- filter(
- lambda task: "dense" in task.task_name,
- extracted_tasks,
- )
- )
- config = ms.TuneConfig(
- strategy="replay_trace",
- num_trials_per_iter=8,
- max_trials_per_task=8,
- max_trials_global=8,
- )
-
- with tempfile.TemporaryDirectory() as work_dir:
- # postprocs=lambda: [] is important to prevent default post processors from
- # tampering with the manual schedule.
- database = ms.tune_extracted_tasks(
- tune_tasks,
- config,
- work_dir=work_dir,
- postprocs=lambda: [],
- )
- else:
-
- def schedule_fn(sch) -> bool:
- if "dense" not in sch.mod.attrs["task_name"]:
- return False
-
- block = sch.get_block("compute")
-
- # Looks up schedule_rule annotation.
- # See the comment in test_tune_relay_manual_tir_vnni().
- schedule_rule = sch.get(block).annotations["schedule_rule"]
-
- assert "dense_vnni" in schedule_rule
-
- schedule_dense(block, M, False, sch)
-
- return True
-
- database = ms.database.ScheduleFnDatabase(schedule_fn)
-
- with database, tvm.transform.PassContext(
- opt_level=3,
- config={"relay.backend.use_meta_schedule": True},
- ):
- # pylint: disable=W0105
- """
- The log should say
- Warning: Cannot find workload: tvmgen_default_fused_expand_dims
- Warning: Cannot find workload: tvmgen_default_fused_cast
- Warning: Cannot find workload: tvmgen_default_fused_cast_1
- Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
-
- This means batch matmul and others are scheduled by TE, and dense (the one not warned)
- is found in the meta schedule tuning database during compilation
- """
- # pylint: enable=W0105
- lib = relay.build(relay_mod, target=target, params=params)
-
- runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
-
- runtime.set_input("data", data)
- runtime.run()
-
- out = runtime.get_output(0).numpy()
-
- np.testing.assert_equal(out, ref)
-
-
-@tvm.testing.requires_cascadelake
-def test_tune_relay_manual_tir_vnni():
Review Comment:
This test and `test_meta_schedule_relay_lowering` above were not ported to new test files.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988351184
##########
python/tvm/meta_schedule/postproc/postproc.py:
##########
@@ -70,6 +74,35 @@ def clone(self) -> "Postproc":
"""
return _ffi_api.PostprocClone(self) # type: ignore # pylint: disable=no-member
+ @staticmethod
+ def create(kind: Literal["llvm", "cuda", "cuda-tensorcore"]) -> List["Postproc"]:
Review Comment:
added
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985191082
##########
src/meta_schedule/task_scheduler/task_scheduler.cc:
##########
@@ -21,83 +21,225 @@
namespace tvm {
namespace meta_schedule {
-void TaskSchedulerNode::InitializeTask(int task_id) {
+TaskRecord::TaskRecord(TuneContext ctx, double task_weight) {
+ ObjectPtr<TaskRecordNode> n = runtime::make_object<TaskRecordNode>();
+ n->ctx = ctx;
+ n->task_weight = task_weight;
+ n->flop = 1.0;
auto _ = Profiler::TimedScope("InitializeTask");
- TuneContext task = this->tasks[task_id];
- TVM_PY_LOG(INFO, this->logging_func)
- << "Initializing Task #" << task_id << ": " << task->task_name;
- TVM_PY_LOG(INFO, task->logging_func)
- << "Initializing Task #" << task_id << ": " << task->task_name;
- CHECK(task->mod.defined()) << "ValueError: Require `context.mod`, but it is not defined";
- CHECK(task->space_generator.defined())
+ CHECK(ctx->mod.defined()) << "ValueError: Require `context.mod`, but it is not defined";
+ CHECK(ctx->space_generator.defined())
<< "ValueError: Require `context.space_generator`, but it is not defined";
- CHECK(task->search_strategy.defined())
+ CHECK(ctx->search_strategy.defined())
<< "ValueError: Require `context.search_strategy`, but it is not defined";
- TVM_PY_LOG(INFO, task->logging_func) << "\n" << tir::AsTVMScript(task->mod);
- task->Initialize();
- Array<tir::Schedule> design_spaces =
- task->space_generator.value()->GenerateDesignSpace(task->mod.value());
- TVM_PY_LOG(INFO, task->logging_func)
- << "Total " << design_spaces.size() << " design space(s) generated";
- for (int i = 0, n = design_spaces.size(); i < n; ++i) {
- tir::Schedule sch = design_spaces[i];
- tir::Trace trace = sch->trace().value();
- trace = trace->Simplified(true);
- TVM_PY_LOG(INFO, task->logging_func) << "Design space #" << i << ":\n"
- << tir::AsTVMScript(sch->mod()) << "\n"
- << Concat(trace->AsPython(false), "\n");
+ TVM_PY_LOG(INFO, ctx->logger) << "\n" << tir::AsTVMScript(ctx->mod);
+ ctx->Initialize();
+ n->flop = std::max(1.0, tir::EstimateTIRFlops(ctx->mod.value()));
+ this->data_ = std::move(n);
+}
+
+void SendToBuilder(TaskRecordNode* self, const Builder& builder) {
+ auto _ = Profiler::TimedScope("SendToBuilder");
+ Array<MeasureCandidate> candidates = self->measure_candidates.value();
+ Target target = self->ctx->target.value();
+ Array<BuilderInput> inputs;
+ inputs.reserve(candidates.size());
+ for (const MeasureCandidate& candidate : candidates) {
+ inputs.push_back(BuilderInput(candidate->sch->mod(), target));
}
- task->search_strategy.value()->PreTuning(design_spaces, database, cost_model);
+ self->builder_results = builder->Build(inputs);
}
-void TaskSchedulerNode::Tune() {
- int n_tasks = this->tasks.size();
- for (int task_id = 0; task_id < n_tasks; ++task_id) {
- InitializeTask(task_id);
+void SendToRunner(TaskRecordNode* self, const Runner& runner) {
+ auto _ = Profiler::TimedScope("SendToRunner");
+ Array<MeasureCandidate> candidates = self->measure_candidates.value();
+ Array<BuilderResult> builder_results = self->builder_results.value();
+ Target target = self->ctx->target.value();
+ ICHECK_EQ(candidates.size(), builder_results.size());
+ int n = candidates.size();
+ int n_build_errors = 0;
+ Array<RunnerInput> inputs;
+ inputs.reserve(n);
+ for (int i = 0; i < n; ++i) {
+ const MeasureCandidate& candidate = candidates[i];
+ const BuilderResult& builder_result = builder_results[i];
+ if (builder_result->error_msg.defined()) {
+ ++n_build_errors;
+ continue;
+ }
+ inputs.push_back(RunnerInput(/*artifact_path=*/builder_result->artifact_path.value(),
+ /*device_type=*/target->kind->name,
+ /*args_info=*/candidate->args_info));
+ }
+ Array<RunnerFuture> futures = runner->Run(inputs);
+ if (n_build_errors == 0) {
+ self->runner_futures = futures;
+ return;
+ }
+ Array<RunnerFuture> results;
+ results.reserve(n);
+ for (int i = 0, j = 0; i < n; ++i) {
+ const BuilderResult& builder_result = builder_results[i];
+ if (builder_result->error_msg.defined()) {
+ results.push_back(RunnerFuture(
+ /*f_done=*/[]() -> bool { return true; },
+ /*f_result=*/
+ [msg = builder_result->error_msg]() -> RunnerResult {
+ return RunnerResult(NullOpt, msg);
+ }));
+ } else {
+ results.push_back(futures[j++]);
+ }
+ }
+ self->runner_futures = results;
+}
+
+void TaskCleanUp(TaskRecordNode* self, int task_id, const Array<RunnerResult>& results) {
+ ICHECK_EQ(self->builder_results.value().size(), results.size());
+ ICHECK_EQ(self->runner_futures.value().size(), results.size());
+ int n = results.size();
+ std::string name = self->ctx->task_name.value();
+ const PackedFunc& logger = self->ctx->logger;
+ for (int i = 0; i < n; ++i) {
+ const BuilderResult& builder_result = self->builder_results.value()[i];
+ const MeasureCandidate& candidate = self->measure_candidates.value()[i];
+ const RunnerResult& runner_result = results[i];
+ Optional<String> error_msg = NullOpt;
+ int trials = self->latency_ms.size() + 1;
+ double run_ms = 1e9;
+ if ((error_msg = builder_result->error_msg)) {
+ ++self->build_error_count;
+ } else if ((error_msg = runner_result->error_msg)) {
+ ++self->run_error_count;
+ } else {
+ run_ms = GetRunMsMedian(runner_result);
+ }
+ self->latency_ms.push_back(run_ms);
+ if (error_msg) {
Review Comment:
It could be pretty frequent if something goes wrong, but normally it didn't happen often. By default, the error message is only sent to a per-task logging file rather than the main screen, so I wouldn't worry too much about it
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190427
##########
include/tvm/meta_schedule/search_strategy.h:
##########
@@ -257,8 +249,8 @@ class PySearchStrategyNode : public SearchStrategyNode {
}
void InitializeWithTuneContext(const TuneContext& context) final;
- void PreTuning(const Array<tir::Schedule>& design_spaces, const Optional<Database>& database,
- const Optional<CostModel>& cost_model) final;
+ void PreTuning(int max_trials, int num_trials_per_iter, const Array<tir::Schedule>& design_spaces,
Review Comment:
It limits the maximum number of trials for a single task. For example, search ends once # trials reaches it: https://github.com/apache/tvm/blob/main/src/meta_schedule/search_strategy/evolutionary_search.cc#L676
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
vinx13 commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r982690510
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -51,6 +51,125 @@ ScheduleRule ScheduleRule::PyScheduleRule(
return ScheduleRule(n);
}
+Array<ScheduleRule> ScheduleRule::DefaultLLVM() {
+ return {
+ ScheduleRule::AutoInline(
+ /*into_producer=*/false,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/true,
+ /*require_injective=*/true,
+ /*require_ordered=*/true,
+ /*disallow_op=*/Array<String>{"tir.exp"}),
+ ScheduleRule::AddRFactor(
+ /*max_jobs_per_core=*/16,
+ /*max_innermost_factor=*/Integer(64)),
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSRSRS",
+ /*tile_binds=*/NullOpt,
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/NullOpt,
+ /*reuse_read=*/NullOpt,
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("may")},
+ {"levels", Array<Integer>{1, 2}},
+ {"scope", String("global")}}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/16,
+ /*max_vectorize_extent=*/64,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512},
+ /*unroll_explicit=*/true),
+ ScheduleRule::RandomComputeLocation(),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
+ return {
+ ScheduleRule::MultiLevelTiling(
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}}),
+ ScheduleRule::AutoInline(
+ /*into_producer=*/true,
+ /*into_consumer=*/true,
+ /*inline_const_tensor=*/true,
+ /*disallow_if_then_else=*/false,
+ /*require_injective=*/false,
+ /*require_ordered=*/false,
+ /*disallow_op=*/Array<String>{}),
+ ScheduleRule::CrossThreadReduction(
+ /*thread_extents=*/Array<Integer>{4, 8, 16, 32, 64, 128, 256, 512}),
+ ScheduleRule::ParallelizeVectorizeUnroll(
+ /*max_jobs_per_core=*/-1,
+ /*max_vectorize_extent=*/-1,
+ /*unroll_max_steps=*/Array<Integer>{0, 16, 64, 512, 1024},
+ /*unroll_explicit=*/true),
+ ScheduleRule::AutoBind(
+ /*max_threadblocks=*/256,
+ /*thread_extents*/ Array<Integer>{32, 64, 128, 256, 512, 1024}),
+ };
+}
+
+Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
+ Array<Map<String, String>> intrin_groups = {
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_f16"},
+ {"load_a", "wmma_load_16x16x16_f16_a"},
+ {"load_b", "wmma_load_16x16x16_f16_b_trans"},
+ {"compute", "wmma_sync_16x16x16_f16f16f16_trans"},
+ {"store", "wmma_store_16x16x16_f16_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ {
+ {"init", "wmma_fill_16x16x16_s32"},
+ {"load_a", "wmma_load_16x16x16_s8_a"},
+ {"load_b", "wmma_load_16x16x16_s8_b_trans"},
+ {"compute", "wmma_sync_16x16x16_s8s8s32_trans"},
+ {"store", "wmma_store_16x16x16_s32_shared"},
+ },
+ };
+ Array<ScheduleRule> results{ScheduleRule::MultiLevelTilingTensorCore(
+ /*intrin_groups=*/intrin_groups,
+ /*structure=*/"SSSRRSRS",
+ /*tile_binds=*/Array<String>{"blockIdx.x", "vthread.x", "threadIdx.x"},
+ /*max_innermost_factor=*/Integer(64),
+ /*vector_load_lens=*/Array<Integer>{1, 2, 3, 4, 8, 16},
+ /*reuse_read=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{4}}, //
+ {"scope", String("shared")}},
+ /*reuse_write=*/
+ Map<String, ObjectRef>{{"req", String("must")},
+ {"levels", Array<Integer>{3}}, //
+ {"scope", String("local")}},
Review Comment:
this is not consistent with previous setting
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
yelite commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r983952886
##########
python/tvm/meta_schedule/utils.py:
##########
@@ -114,7 +113,10 @@ def __init__(self, *args, **kwargs):
def __getattr__(self, name: str):
"""Bridge the attribute function."""
- return self._inst.__getattribute__(name)
+ try:
+ return self._inst.__getattribute__(name)
+ except AttributeError:
+ return super(TVMDerivedObject, self).__getattr__(name)
Review Comment:
Not really related to the change here. But is it possible to make mypy treat DerivedObject as sub class of the real base class from `_tvm_metadata`? For example, right now `RPCRunner` is not a subclass of `Runner` from mypy's perspective.
##########
python/tvm/meta_schedule/logging.py:
##########
@@ -0,0 +1,262 @@
+# 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.
+"""Logging interface in MetaSchedule"""
+import logging
+import logging.config
+import os
+import os.path as osp
+from logging import Logger
+from typing import Any, Callable, Dict, List, Optional
+
+
+def get_logger(name: str) -> Logger:
+ """Create or get a logger by its name. This is essentially a wrapper of python's native logger.
+
+ Parameters
+ ----------
+ name : str
+ The name of the logger.
+
+ Returns
+ -------
+ logger : Logger
+ The logger instance.
+ """
+ return logging.getLogger(name)
+
+
+def make_logging_func(logger: Logger) -> Optional[Callable[[int, str], None]]:
+ """Get the logging function.
+
+ Parameters
+ ----------
+ logger : Logger
+ The logger instance.
+ Returns
+ -------
+ result : Optional[Callable]
+ The function to do the specified level of logging.
+ """
+ if logger is None:
+ return None
+
+ level2log = {
+ logging.DEBUG: logger.debug,
+ logging.INFO: logger.info,
+ logging.WARNING: logger.warning,
+ logging.ERROR: logger.error,
+ # logging.FATAL not included
+ }
+
+ def logging_func(level: int, msg: str):
+ def clear_notebook_output():
+ from IPython.display import ( # type: ignore # pylint: disable=import-outside-toplevel
+ clear_output,
+ )
+
+ clear_output(wait=True)
+
+ if level < 0:
+ clear_notebook_output()
Review Comment:
Will there be the case that `level < 0` while code isn't executed in notebook?
##########
python/tvm/meta_schedule/logging.py:
##########
@@ -0,0 +1,262 @@
+# 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.
+"""Logging interface in MetaSchedule"""
+import logging
+import logging.config
+import os
+import os.path as osp
+from logging import Logger
+from typing import Any, Callable, Dict, List, Optional
+
+
+def get_logger(name: str) -> Logger:
+ """Create or get a logger by its name. This is essentially a wrapper of python's native logger.
+
+ Parameters
+ ----------
+ name : str
+ The name of the logger.
+
+ Returns
+ -------
+ logger : Logger
+ The logger instance.
+ """
+ return logging.getLogger(name)
+
+
+def make_logging_func(logger: Logger) -> Optional[Callable[[int, str], None]]:
+ """Get the logging function.
+
+ Parameters
+ ----------
+ logger : Logger
+ The logger instance.
+ Returns
+ -------
+ result : Optional[Callable]
+ The function to do the specified level of logging.
+ """
+ if logger is None:
+ return None
+
+ level2log = {
+ logging.DEBUG: logger.debug,
+ logging.INFO: logger.info,
+ logging.WARNING: logger.warning,
+ logging.ERROR: logger.error,
+ # logging.FATAL not included
+ }
+
+ def logging_func(level: int, msg: str):
+ def clear_notebook_output():
+ from IPython.display import ( # type: ignore # pylint: disable=import-outside-toplevel
+ clear_output,
+ )
+
+ clear_output(wait=True)
+
+ if level < 0:
+ clear_notebook_output()
Review Comment:
Will this be better?
```suggestion
if level < 0:
from IPython.display import clear_output
clear_output(wait=True)
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985191319
##########
include/tvm/meta_schedule/search_strategy.h:
##########
@@ -88,14 +88,17 @@ class SearchStrategyNode : public runtime::Object {
/*!
* \brief Pre-tuning for the search strategy.
+ * \param max_trials The maximum number of trials.
+ * \param num_trials_per_iter The number of trials per iteration.
* \param design_spaces The design spaces used during tuning process.
Review Comment:
in fact, it's pretty consistent in the codebase - when it's a list, we add an `s`
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r990227657
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
Review Comment:
this should work
```
space=ms.space_generator.PostOrderApply(
f_block_filter=None,
sch_rules="from-target",
postprocs=[],
mutator_probs="from-target",
),
)
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985051195
##########
tests/python/unittest/test_meta_schedule_space_generator.py:
##########
@@ -104,4 +108,5 @@ class TestPySpaceGenerator(PySpaceGenerator):
if __name__ == "__main__":
- tvm.testing.main()
+ test_meta_schedule_design_space_generator_NIE()
+ # tvm.testing.main()
Review Comment:
oh ooops slipped it through
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988351447
##########
src/meta_schedule/postproc/postproc.cc:
##########
@@ -67,6 +109,10 @@ TVM_REGISTER_GLOBAL("meta_schedule.PostprocInitializeWithTuneContext")
TVM_REGISTER_GLOBAL("meta_schedule.PostprocApply").set_body_method<Postproc>(&PostprocNode::Apply);
TVM_REGISTER_GLOBAL("meta_schedule.PostprocClone").set_body_method<Postproc>(&PostprocNode::Clone);
TVM_REGISTER_GLOBAL("meta_schedule.PostprocPyPostproc").set_body_typed(Postproc::PyPostproc);
+TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultLLVM").set_body_typed(Postproc::DefaultLLVM);
+TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultCUDA").set_body_typed(Postproc::DefaultCUDA);
+TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultCUDATensorCore")
+ .set_body_typed(Postproc::DefaultCUDATensorCore);
Review Comment:
added
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] csullivan commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
csullivan commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r1002899944
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
+ # pylint: disable=W0105
+ """
+ We can inject and apply a custom TIR scheduling to a TE compute of interest, using
+ the "schedule_rule" annotation. For example, in topi/x86/dense.py we have the following
+ declaration for int8 dense targeting the VNNI instruction.
+
+ C = te.compute(
+ ...
+ attrs={"schedule_rule": "meta_schedule.dense_vnni"},
+ )
+
+ When the MetaSchedule encounters a TensorIR block with the "schedule_rule" annotation,
+ it looks up the packed func registry for a function that is associated with the given schedule
+ rule key ("meta_schedule.dense_vnni" in this example). The signature of such custom schedule
+ functions must be
+
+ (tir.schedule.Schedule, tir.schedule.BlockRV) -> [tir.schedule.Schedule].
+
+ The BlockRV argument corresponds to the TE compute annotated with "schedule_rule".
+
+ The relevant code is in meta_schedule/space_generator/post_order_apply.cc.
+ """
+
+ def schedule_rule_dense_vnni(sch: Schedule, dense_block: BlockRV):
+ _schedule_dense(m=None, do_tune=True)(sch, dense_block)
+ return [sch]
+
+ register_func("meta_schedule.dense_vnni", schedule_rule_dense_vnni)
Review Comment:
@junrushao @masahi or others, may I ask what the difference is between using the TE annotation as described (e.g. `attrs={"schedule_rule": "meta_schedule.dense_vnni"},` and a corresponding packed func defining the schedule to use, as opposed to just generating the space via
```
space=ms.space_generator.ScheduleFn(
_schedule_dense,
...
),
```
?
Is it that in this test case we allow auto scheduling for all ops but apply special manual scheduling for certain ops (dense in this case), whereas if we use the ScheduleFn technique for generating a search space we do not allow other operators to be auto scheduled? Thanks!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r1002954336
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
+ # pylint: disable=W0105
+ """
+ We can inject and apply a custom TIR scheduling to a TE compute of interest, using
+ the "schedule_rule" annotation. For example, in topi/x86/dense.py we have the following
+ declaration for int8 dense targeting the VNNI instruction.
+
+ C = te.compute(
+ ...
+ attrs={"schedule_rule": "meta_schedule.dense_vnni"},
+ )
+
+ When the MetaSchedule encounters a TensorIR block with the "schedule_rule" annotation,
+ it looks up the packed func registry for a function that is associated with the given schedule
+ rule key ("meta_schedule.dense_vnni" in this example). The signature of such custom schedule
+ functions must be
+
+ (tir.schedule.Schedule, tir.schedule.BlockRV) -> [tir.schedule.Schedule].
+
+ The BlockRV argument corresponds to the TE compute annotated with "schedule_rule".
+
+ The relevant code is in meta_schedule/space_generator/post_order_apply.cc.
+ """
+
+ def schedule_rule_dense_vnni(sch: Schedule, dense_block: BlockRV):
+ _schedule_dense(m=None, do_tune=True)(sch, dense_block)
+ return [sch]
+
+ register_func("meta_schedule.dense_vnni", schedule_rule_dense_vnni)
Review Comment:
I think `ScheduleFnDatabase` is for a completely manual schedule, while the `register_func` way allows autotvm-style template based tuning. At least that's what I wanted to demonstrate before this PR or before `ScheduleFnDatabase` was introduced.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r984244564
##########
include/tvm/meta_schedule/search_strategy.h:
##########
@@ -88,14 +88,17 @@ class SearchStrategyNode : public runtime::Object {
/*!
* \brief Pre-tuning for the search strategy.
+ * \param max_trials The maximum number of trials.
+ * \param num_trials_per_iter The number of trials per iteration.
* \param design_spaces The design spaces used during tuning process.
Review Comment:
I noticed that we use plural here for `design_spaces`, I kind of prefer to use singular term for consistency, but no strong opinion.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988350143
##########
src/meta_schedule/postproc/postproc.cc:
##########
@@ -50,6 +50,48 @@ Postproc Postproc::PyPostproc(
return Postproc(n);
}
+Array<Postproc> Postproc::DefaultLLVM() {
+ return Array<Postproc>{
+ Postproc::DisallowDynamicLoop(),
+ Postproc::RewriteParallelVectorizeUnroll(),
+ Postproc::RewriteReductionBlock(),
+ Postproc::RewriteLayout(),
+ };
+}
+
+Array<Postproc> Postproc::DefaultCUDA() {
+ return Array<Postproc>{
+ Postproc::DisallowDynamicLoop(),
+ Postproc::RewriteCooperativeFetch(),
+ Postproc::RewriteUnboundBlock(/*max_threadblocks=*/256),
+ Postproc::RewriteParallelVectorizeUnroll(),
+ Postproc::RewriteReductionBlock(),
+ Postproc::VerifyGPUCode(),
+ };
+}
+
+Array<Postproc> Postproc::DefaultCUDATensorCore() {
+ return Array<Postproc>{
+ Postproc::DisallowDynamicLoop(),
+ Postproc::RewriteCooperativeFetch(),
+ Postproc::RewriteUnboundBlock(/*max_threadblocks=*/256),
+ Postproc::RewriteParallelVectorizeUnroll(),
+ Postproc::RewriteReductionBlock(),
+ Postproc::RewriteTensorize(/*vectorize_init_loop=*/false),
+ Postproc::VerifyGPUCode(),
+ };
+}
+
+Array<Postproc> Postproc::DefaultHexagon() {
+ return Array<Postproc>{
+ Postproc::DisallowDynamicLoop(),
+ Postproc::RewriteParallelVectorizeUnroll(), //
+ Postproc::RewriteReductionBlock(),
+ // TODO(masahi): Fix RewriteLayout for link-params=True case
+ // Postproc::RewriteLayout(),
+ };
+}
Review Comment:
got it. done!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989378918
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,101 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+if TYPE_CHECKING:
+ from tvm import relay
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=True,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+ pass_config: Mapping[str, Any],
+ executor: Optional["relay.backend.Executor"],
+) -> Tuple[IRModule, Target, Dict[str, NDArray], Dict[str, Any], "relay.backend.Executor"]:
+ from tvm import relay # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, relay.Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+ if executor is None:
+ executor = relay.backend.Executor("graph")
+ pass_config = dict(pass_config)
+ pass_config.setdefault(
+ "relay.FuseOps.link_params",
+ executor.attrs.get("link_params", False),
+ )
Review Comment:
oh i see the logic here: https://github.com/apache/tvm/blob/d4bf9ecf5524d265916ac7b860b0027f5eee5c49/src/relay/transforms/fuse_ops.cc#L1052-L1057
okay will do!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r989636824
##########
tests/python/contrib/test_hexagon/test_meta_schedule.py:
##########
@@ -358,6 +361,7 @@ def test_conv2d_relay_auto_schedule(hexagon_launcher):
kernel_layout="HWIO",
)
mod = tvm.IRModule.from_expr(conv2d + bias)
+ mod = mod.with_attr("executor", relay.backend.Executor("graph"))
Review Comment:
Oh got it! It somehow slipped through
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1270853843
@masahi I updated the PR with my latest understanding of Hexagon pipeline. Would you mind taking another look? Thanks a lot!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988341530
##########
python/tvm/meta_schedule/cost_model/cost_model.py:
##########
@@ -97,6 +104,30 @@ def predict(self, context: TuneContext, candidates: List[MeasureCandidate]) -> n
)
return results
+ @staticmethod
+ def create(
+ kind: Literal["xgb", "mlp", "random"],
Review Comment:
thanks for bringing this up!
##########
python/tvm/meta_schedule/cost_model/cost_model.py:
##########
@@ -97,6 +104,30 @@ def predict(self, context: TuneContext, candidates: List[MeasureCandidate]) -> n
)
return results
+ @staticmethod
+ def create(
+ kind: Literal["xgb", "mlp", "random"],
Review Comment:
done
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] csullivan commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
csullivan commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r1003440566
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
+ # pylint: disable=W0105
+ """
+ We can inject and apply a custom TIR scheduling to a TE compute of interest, using
+ the "schedule_rule" annotation. For example, in topi/x86/dense.py we have the following
+ declaration for int8 dense targeting the VNNI instruction.
+
+ C = te.compute(
+ ...
+ attrs={"schedule_rule": "meta_schedule.dense_vnni"},
+ )
+
+ When the MetaSchedule encounters a TensorIR block with the "schedule_rule" annotation,
+ it looks up the packed func registry for a function that is associated with the given schedule
+ rule key ("meta_schedule.dense_vnni" in this example). The signature of such custom schedule
+ functions must be
+
+ (tir.schedule.Schedule, tir.schedule.BlockRV) -> [tir.schedule.Schedule].
+
+ The BlockRV argument corresponds to the TE compute annotated with "schedule_rule".
+
+ The relevant code is in meta_schedule/space_generator/post_order_apply.cc.
+ """
+
+ def schedule_rule_dense_vnni(sch: Schedule, dense_block: BlockRV):
+ _schedule_dense(m=None, do_tune=True)(sch, dense_block)
+ return [sch]
+
+ register_func("meta_schedule.dense_vnni", schedule_rule_dense_vnni)
Review Comment:
In this case I'm not referring to ScheduleFnDatabase as is used in `test_vnni_schedule_fn_database`. I'm referring here to what is done in the test `test_vnni_schedule_fn_tune` which utilizes the TE compute `schedule_rule` attr annotation along with a global packed function for the schedule that matches the annotation value `meta_schedule.dense_vnni`. I'm wondering if there is any difference or advantage between using the TE attr annotation and packed func as opposed to specifying an alternate search space with ScheduleFn.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r990224380
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
Review Comment:
will send a fix
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1260475610
@tqchen @Hzfengsy @spectrometerHBH @zxybazh @vinx13 @yelite The PR is ready for review. Please take a look!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r1003609587
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
+ # pylint: disable=W0105
+ """
+ We can inject and apply a custom TIR scheduling to a TE compute of interest, using
+ the "schedule_rule" annotation. For example, in topi/x86/dense.py we have the following
+ declaration for int8 dense targeting the VNNI instruction.
+
+ C = te.compute(
+ ...
+ attrs={"schedule_rule": "meta_schedule.dense_vnni"},
+ )
+
+ When the MetaSchedule encounters a TensorIR block with the "schedule_rule" annotation,
+ it looks up the packed func registry for a function that is associated with the given schedule
+ rule key ("meta_schedule.dense_vnni" in this example). The signature of such custom schedule
+ functions must be
+
+ (tir.schedule.Schedule, tir.schedule.BlockRV) -> [tir.schedule.Schedule].
+
+ The BlockRV argument corresponds to the TE compute annotated with "schedule_rule".
+
+ The relevant code is in meta_schedule/space_generator/post_order_apply.cc.
+ """
+
+ def schedule_rule_dense_vnni(sch: Schedule, dense_block: BlockRV):
+ _schedule_dense(m=None, do_tune=True)(sch, dense_block)
+ return [sch]
+
+ register_func("meta_schedule.dense_vnni", schedule_rule_dense_vnni)
Review Comment:
Hi Chris, `ScheduleFn` space generator is designed to schedule all blocks in the whole Schedule, not block specific. The annotation based packfunc scheduling only works in `PostOrderApply` space generator, which essentially applies this annotated rule for this specific block, and apply default schedule rules (or schedule rules given in user interface) to other non-specified blocks.
Therefore, creating a `ScheduleFn` takes more effort and use the annotation based scheduling is easier because you don't need to worry about scheduling of other blocks.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] csullivan commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
csullivan commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r1003817407
##########
tests/python/unittest/test_meta_schedule_vnni_integration.py:
##########
@@ -0,0 +1,249 @@
+# 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-docstring
+import logging
+import tempfile
+from typing import Optional
+
+import numpy as np # type: ignore
+import pytest
+import tvm
+from tvm import meta_schedule as ms
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.tir.schedule import BlockRV, Schedule
+from tvm.tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
+
+logging.basicConfig(
+ format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+logging.getLogger("tvm.meta_schedule").setLevel(logging.DEBUG)
+
+
+def _schedule_dense(m: Optional[int], do_tune: bool):
+ """Manually schedule a dense block, created from TE compute op via CreatePrimFunc,
+ using VNNI instruction.
+ """
+
+ def schedule_fn(sch, dense_block: Optional[BlockRV] = None) -> bool:
+ if "dense" not in sch.mod.attrs["task_name"]:
+ return False
+ if dense_block is None:
+ dense_block = sch.get_block("compute")
+ assert "dense_vnni" in sch.get(dense_block).annotations["schedule_rule"]
+
+ post_blocks = sch.get_consumers(dense_block)
+ if len(post_blocks) > 0:
+ # Fuse all intermediate post ops into the last op.
+ # This is equivalent to the traverse_inline function used in TE schedules.
+ while True:
+ next_post_blocks = []
+ for post_block in post_blocks:
+ next_consumers = sch.get_consumers(post_block)
+ if len(next_consumers) > 0:
+ sch.compute_inline(post_block)
+ next_post_blocks += next_consumers
+ if len(next_post_blocks) == 0:
+ assert len(post_blocks) == 1
+ outer_block = post_blocks[0]
+ a_y, a_x = sch.get_loops(outer_block)[-2:]
+ break
+ post_blocks = next_post_blocks
+ else:
+ a_y, a_x, _ = sch.get_loops(dense_block)[-3:]
+ outer_block = dense_block
+ if do_tune:
+ y_factors = sch.sample_perfect_tile(a_y, n=2, max_innermost_factor=128)
+ a_yo, a_yi = sch.split(a_y, factors=y_factors)
+ else:
+ a_yo, a_yi = sch.split(a_y, factors=[None, min(m, 64)])
+ a_xo, a_xi = sch.split(a_x, factors=[None, 16])
+ sch.reorder(a_yo, a_xo, a_yi, a_xi)
+ fused = sch.fuse(a_yo, a_xo)
+ if outer_block != dense_block:
+ # Handle the case when dense is fused with post ops.
+ sch.vectorize(a_xi)
+ sch.compute_at(dense_block, a_yi)
+ a_xi, a_k = sch.get_loops(dense_block)[-2:]
+ a_ko, a_ki = sch.split(a_k, factors=[None, 4])
+ sch.reorder(a_ko, a_xi, a_ki)
+ # We need to parallelize before decompose_reduction, otherwise the so-called "Compact dataflow"
+ # condition is violated.
+ sch.parallel(fused)
+ dec = sch.decompose_reduction(dense_block, a_ko)
+ init_loop = sch.get_loops(dec)[-1]
+ sch.vectorize(init_loop)
+ sch.tensorize(a_xi, VNNI_INTRIN)
+ return True
+
+ return schedule_fn
+
+
+def _relay_dense(m, n, k):
+ data = relay.var("data", shape=(m, k), dtype="uint8")
+ weight = relay.var("weight", shape=(n, k), dtype="int8")
+ bias = relay.var("bias", shape=(n,), dtype="int32")
+ # dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
+ dense = relay.nn.dense(data, weight, out_dtype="int32")
+ bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
+ out = relay.nn.batch_matmul(
+ relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
+ relay.cast(relay.expand_dims(bias_add, 0), "int8"),
+ out_dtype="int32",
+ )
+ relay_mod = tvm.IRModule.from_expr(out)
+ data = np.random.uniform(1, 10, size=(m, k)).astype("uint8")
+ params = {
+ "weight": np.random.uniform(1, 10, size=(n, k)).astype("int8"),
+ "bias": np.random.uniform(1, 10, size=(n,)).astype("int32"),
+ }
+
+ def f_check(lib, dev):
+ ref = (
+ relay.create_executor(
+ "vm",
+ mod=relay_mod,
+ device=dev,
+ target="llvm",
+ )
+ .evaluate()(data, params["weight"], params["bias"])
+ .numpy()
+ )
+ runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ runtime.set_input("data", data)
+ runtime.run()
+ out = runtime.get_output(0).numpy()
+ np.testing.assert_equal(out, ref)
+
+ return relay_mod, params, f_check
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_database():
+ m, n, k = 1024, 1024, 1024
+ target = tvm.target.Target("llvm -mcpu=cascadelake -num-cores 4")
+ dev = tvm.cpu(0)
+ relay_mod, params, f_check = _relay_dense(m, n, k)
+
+ with ms.database.ScheduleFnDatabase(
+ _schedule_dense(
+ m=m,
+ do_tune=False,
+ )
+ ), tvm.transform.PassContext(
+ opt_level=3,
+ config={"relay.backend.use_meta_schedule": True},
+ ):
+ # pylint: disable=W0105
+ """The log should say
+ Warning: Cannot find workload: tvmgen_default_fused_expand_dims
+ Warning: Cannot find workload: tvmgen_default_fused_cast
+ Warning: Cannot find workload: tvmgen_default_fused_cast_1
+ Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
+
+ This means batch matmul and others are scheduled by TE, and dense (the one not warned)
+ is found in the meta schedule tuning database during compilation
+ """
+ # pylint: enable=W0105
+ lib = relay.build(relay_mod, target=target, params=params)
+ f_check(lib, dev)
+
+
+@pytest.mark.skip("Requires cascadelake")
+def test_vnni_schedule_fn_tune():
+ # pylint: disable=W0105
+ """
+ We can inject and apply a custom TIR scheduling to a TE compute of interest, using
+ the "schedule_rule" annotation. For example, in topi/x86/dense.py we have the following
+ declaration for int8 dense targeting the VNNI instruction.
+
+ C = te.compute(
+ ...
+ attrs={"schedule_rule": "meta_schedule.dense_vnni"},
+ )
+
+ When the MetaSchedule encounters a TensorIR block with the "schedule_rule" annotation,
+ it looks up the packed func registry for a function that is associated with the given schedule
+ rule key ("meta_schedule.dense_vnni" in this example). The signature of such custom schedule
+ functions must be
+
+ (tir.schedule.Schedule, tir.schedule.BlockRV) -> [tir.schedule.Schedule].
+
+ The BlockRV argument corresponds to the TE compute annotated with "schedule_rule".
+
+ The relevant code is in meta_schedule/space_generator/post_order_apply.cc.
+ """
+
+ def schedule_rule_dense_vnni(sch: Schedule, dense_block: BlockRV):
+ _schedule_dense(m=None, do_tune=True)(sch, dense_block)
+ return [sch]
+
+ register_func("meta_schedule.dense_vnni", schedule_rule_dense_vnni)
Review Comment:
Ahh, okay I see, thanks for the discussion @zxybazh @masahi, this is helpful.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r983983542
##########
python/tvm/meta_schedule/utils.py:
##########
@@ -114,7 +113,10 @@ def __init__(self, *args, **kwargs):
def __getattr__(self, name: str):
"""Bridge the attribute function."""
- return self._inst.__getattribute__(name)
+ try:
+ return self._inst.__getattribute__(name)
+ except AttributeError:
+ return super(TVMDerivedObject, self).__getattr__(name)
Review Comment:
[Updated link] I think it's a [known issue](https://github.com/python/mypy/issues/3135) of mypy, which haven't been fixed. I didn't see better solution than mark as ignore though.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190154
##########
python/tvm/meta_schedule/utils.py:
##########
@@ -114,7 +113,10 @@ def __init__(self, *args, **kwargs):
def __getattr__(self, name: str):
"""Bridge the attribute function."""
- return self._inst.__getattribute__(name)
+ try:
+ return self._inst.__getattribute__(name)
+ except AttributeError:
+ return super(TVMDerivedObject, self).__getattr__(name)
Review Comment:
Yeah I was troubled by this issue too, but didn't end up finding a good solution. Given it's a known issue in mypy, probably the best thing we could do is to wait and see if it's going to be fixed in next releases.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988320000
##########
python/tvm/contrib/torch/optimize_torch.py:
##########
@@ -62,75 +51,32 @@ def forward(self, *torch_inputs: Tuple[torch.Tensor]):
return ret
-def llvm_target():
- return "llvm -num-cores"
-
-
@register_func("script_torch.save_to_base64")
def save_to_base64(obj) -> bytes:
with tempfile.NamedTemporaryFile(suffix=".so") as tmpfile:
obj.export_library(tmpfile.name)
- with open(tmpfile.name, "rb") as tfile:
- return base64.b64encode(tfile.read())
-
-
-def tune_relay_auto(
- mod: IRModule,
- target: Union[str, Target],
- config: TuneConfig,
- work_dir: str,
- backend: str = "graph",
- params: Optional[Dict[str, NDArray]] = None,
-) -> Union[Module, vm.Executable]:
- """A wrapper of `tune_relay` but provide a default setting for the config.
-
- Parameters
- ----------
- mod : IRModule
- The module to tune.
- target : Union[str, Target]
- The target to tune for.
- config : TuneConfig
- The search strategy config.
- params : Optional[Dict[str, tvm.runtime.NDArray]]
- The associated parameters of the program
- work_dir : Optional[str]
- The working directory to save intermediate results.
- backend : str = "graph"
- The backend to use for relay compilation(graph / vm).
-
- Returns
- -------
- lib : Union[Module, tvm.runtime.vm.Executable]
- The built runtime module or vm Executable for the given relay workload.
- """
- target = default_config.target(target)
- extracted_tasks = extract_task_from_relay(mod, target, params)
- if config is None:
- config = TuneConfig(
- num_trials_per_iter=16,
- max_trials_global=16 * len(extracted_tasks),
- )
- database = tune_extracted_tasks(extracted_tasks, config, work_dir)
- relay_build = {"graph": relay.build, "vm": relay.vm.compile}[backend]
- with target, autotvm_silencer(), database:
- with PassContext(
- opt_level=3,
- config={
- "relay.backend.use_meta_schedule": True,
- "relay.backend.use_meta_schedule_dispatch": target.kind.name != "cuda",
- "relay.backend.tir_converter": "default",
- },
- ):
- return relay_build(mod, target=target, params=params)
+ with open(tmpfile.name, "rb") as temp_file:
+ return base64.b64encode(temp_file.read())
def optimize_torch(
func,
example_inputs,
- tuning_config=None,
- target=None,
+ *,
work_dir=None,
+ target: Union[str, Target] = "cpu",
+ max_trials_global: int,
+ max_trials_per_task: Optional[int] = None,
+ num_trials_per_iter: int = 64,
+ builder: ms.Builder.BuilderType = "local",
+ runner: ms.Runner.RunnerType = "local",
+ database: ms.Database.DatabaseType = "json",
+ cost_model: ms.CostModel.CostModelType = "xgb",
+ measure_callbacks: ms.MeasureCallback.CallbackListType = "default",
+ task_scheduler: ms.TaskScheduler.TaskSchedulerType = "gradient",
+ space: ms.SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+ strategy: ms.SearchStrategy.SearchStrategyType = "evolutionary",
Review Comment:
sorry i forgot! thanks for pointing it out!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] spectrometerHBH merged pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
spectrometerHBH merged PR #12895:
URL: https://github.com/apache/tvm/pull/12895
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] sunggg commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
sunggg commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985956777
##########
python/tvm/meta_schedule/database/database.py:
##########
@@ -361,6 +363,44 @@ def current() -> Optional["Database"]:
"""Get the current database under scope."""
return _ffi_api.DatabaseCurrent() # type: ignore # pylint: disable=no-member
+ @staticmethod
+ def create( # pylint: disable=keyword-arg-before-vararg
Review Comment:
Just for a curiosity.
Is the motivation for this API to remove user-side imports for various database types?
For example, with this, users wouldn't need to import every database. Rather, they can just import this API.
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -69,47 +117,218 @@ def extract_task_from_relay(
"""
# pylint: disable=import-outside-toplevel
from tvm import autotvm
- from tvm.relay import Function as RelayFunc
# pylint: enable=import-outside-toplevel
+ mod, target, params = _normalize_params(mod, target, params)
+ if target.kind.name != "cuda" and isinstance(
+ autotvm.DispatchContext.current, autotvm.FallbackContext
+ ):
+ tophub_context = autotvm.tophub.context(target)
+ else:
+ tophub_context = autotvm.utils.EmptyContext()
+ with Profiler.timeit("TaskExtraction"):
+ with target, _autotvm_silencer(), tophub_context:
+ with transform.PassContext(
+ opt_level=3,
+ config={
+ "relay.backend.use_meta_schedule": True,
+ "relay.backend.tir_converter": tir_converter,
+ },
+ ):
+ return list(_extract_task(mod, target, params))
- extract_task_func = get_global_func(
- "relay.backend.MetaScheduleExtractTask",
- allow_missing=False,
- )
- if isinstance(mod, RelayFunc):
- mod = IRModule.from_expr(mod)
- if not isinstance(target, Target):
- target = Target(target)
- if disabled_pass is None:
- disabled_pass = []
- if pass_config is None:
- pass_config = {
- "relay.backend.use_meta_schedule": True,
- "relay.backend.tir_converter": tir_converter,
- }
- if params is None:
- params = {}
- relay_params = {}
- for name, param in params.items():
- if isinstance(param, np.ndarray):
- param = nd.array(param)
- relay_params[name] = param
+def extracted_tasks_to_tune_contexts(
+ extracted_tasks: List[ExtractedTask],
+ work_dir: str,
+ space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+ strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+ num_threads: Union[Literal["physical", "logical"], int] = "physical",
+ seed: Optional[int] = None,
+) -> Tuple[List[TuneContext], List[float]]:
+ """Convert ExtractedTask to TuneContext.
- with target, autotvm_silencer(), transform.PassContext(
- opt_level=opt_level,
- config=pass_config,
- disabled_pass=disabled_pass,
+ Parameters
+ ----------
+ tasks : List[ExtractedTask]
+ The tasks to be converted
+ work_dir : str
+ The working directory to store logs and databases
+ space : SpaceGenerator.SpaceGeneratorType
+ The space generator to use.
+ strategy : SearchStrategy.SearchStrategyType
+ The search strategy to use.
+ num_threads : Union[Literal["physical", "logical"], int]
+ The number of threads to use.
+ seed : Optional[int]
+ The random seed to use.
+
+ Returns
+ -------
+ tasks : List[TuneContext]
+ The converted tasks
+ task_weights : List[float]
+ The weights of the tasks
+ """
+ tasks: List[TuneContext] = []
+ task_weights: List[float] = []
+ for task, logger, rand_state in zip(
+ extracted_tasks,
+ get_loggers_from_work_dir(work_dir, [t.task_name for t in extracted_tasks]),
+ fork_seed(seed, n=len(extracted_tasks)),
):
- if target.kind.name != "cuda" and isinstance(
- autotvm.DispatchContext.current, autotvm.FallbackContext
- ):
- tophub_context = autotvm.tophub.context(target)
- else:
- tophub_context = autotvm.utils.EmptyContext()
- with tophub_context:
- return list(extract_task_func(mod, target, relay_params))
+ tasks.append(
+ TuneContext(
+ mod=task.dispatched[0],
+ target=task.target,
+ space_generator=space,
+ search_strategy=strategy,
+ task_name=task.task_name,
+ logger=logger,
+ rand_state=rand_state,
+ num_threads=num_threads,
+ ).clone()
+ )
+ task_weights.append(task.weight)
+ return tasks, task_weights
+
+
+def tune_relay(
Review Comment:
Given that this function calls several functions internally, I think it would be great if we can include their all or important arguments. For example, current function cannot allow customization of `num_threads` for `extracted_tasks_to_tune_contexts`.
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -69,47 +117,218 @@ def extract_task_from_relay(
"""
# pylint: disable=import-outside-toplevel
from tvm import autotvm
- from tvm.relay import Function as RelayFunc
# pylint: enable=import-outside-toplevel
+ mod, target, params = _normalize_params(mod, target, params)
+ if target.kind.name != "cuda" and isinstance(
+ autotvm.DispatchContext.current, autotvm.FallbackContext
+ ):
+ tophub_context = autotvm.tophub.context(target)
+ else:
+ tophub_context = autotvm.utils.EmptyContext()
+ with Profiler.timeit("TaskExtraction"):
+ with target, _autotvm_silencer(), tophub_context:
+ with transform.PassContext(
+ opt_level=3,
+ config={
+ "relay.backend.use_meta_schedule": True,
+ "relay.backend.tir_converter": tir_converter,
+ },
+ ):
+ return list(_extract_task(mod, target, params))
- extract_task_func = get_global_func(
- "relay.backend.MetaScheduleExtractTask",
- allow_missing=False,
- )
- if isinstance(mod, RelayFunc):
- mod = IRModule.from_expr(mod)
- if not isinstance(target, Target):
- target = Target(target)
- if disabled_pass is None:
- disabled_pass = []
- if pass_config is None:
- pass_config = {
- "relay.backend.use_meta_schedule": True,
- "relay.backend.tir_converter": tir_converter,
- }
- if params is None:
- params = {}
- relay_params = {}
- for name, param in params.items():
- if isinstance(param, np.ndarray):
- param = nd.array(param)
- relay_params[name] = param
+def extracted_tasks_to_tune_contexts(
+ extracted_tasks: List[ExtractedTask],
+ work_dir: str,
+ space: SpaceGenerator.SpaceGeneratorType = "post-order-apply",
+ strategy: SearchStrategy.SearchStrategyType = "evolutionary",
+ num_threads: Union[Literal["physical", "logical"], int] = "physical",
+ seed: Optional[int] = None,
+) -> Tuple[List[TuneContext], List[float]]:
+ """Convert ExtractedTask to TuneContext.
- with target, autotvm_silencer(), transform.PassContext(
- opt_level=opt_level,
- config=pass_config,
- disabled_pass=disabled_pass,
+ Parameters
+ ----------
+ tasks : List[ExtractedTask]
+ The tasks to be converted
+ work_dir : str
+ The working directory to store logs and databases
+ space : SpaceGenerator.SpaceGeneratorType
+ The space generator to use.
+ strategy : SearchStrategy.SearchStrategyType
+ The search strategy to use.
+ num_threads : Union[Literal["physical", "logical"], int]
Review Comment:
Is this `num_threads` to use to speed up the compilation? would you clarify this in the comment?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on PR #12895:
URL: https://github.com/apache/tvm/pull/12895#issuecomment-1269457217
> On the other hand, as a high-level API, I would prefer not to tweak tune_relay adding more parameters to it, given we wanted to give a cleaner interface to introductory level users. Instead, advance users could always use extract_tasks + tune_tasks + compile_relay to get fine-grained control over the tuning process
Yes, I agree with this. A part of the reason I didn't want to change `extract_task` API before was that the vast majority of users don't need to care about `executor` stuff. For Hexagon users, we can add a wrapper API in `contrib/hexagon/meta_schedule` to simplify the usage. We already require using Hexagon-specific builder / runner, so the wrapper API can hide such details too.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988351862
##########
src/meta_schedule/mutator/mutator.cc:
##########
@@ -72,6 +97,10 @@ TVM_REGISTER_GLOBAL("meta_schedule.MutatorApply")
});
TVM_REGISTER_GLOBAL("meta_schedule.MutatorClone").set_body_method<Mutator>(&MutatorNode::Clone);
TVM_REGISTER_GLOBAL("meta_schedule.MutatorPyMutator").set_body_typed(Mutator::PyMutator);
+TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultLLVM").set_body_typed(Mutator::DefaultLLVM);
+TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultCUDA").set_body_typed(Mutator::DefaultCUDA);
+TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultCUDATensorCore")
Review Comment:
added
##########
src/meta_schedule/schedule_rule/schedule_rule.cc:
##########
@@ -71,6 +217,12 @@ TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleClone")
.set_body_method<ScheduleRule>(&ScheduleRuleNode::Clone);
TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRulePyScheduleRule")
.set_body_typed(ScheduleRule::PyScheduleRule);
+TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultLLVM")
+ .set_body_typed(ScheduleRule::DefaultLLVM);
+TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultCUDA")
+ .set_body_typed(ScheduleRule::DefaultCUDA);
+TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultCUDATensorCore")
+ .set_body_typed(ScheduleRule::DefaultCUDATensorCore);
Review Comment:
added
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #12895: [MetaSchedule] UX: Tuning API cleanup & developer ergonomics
Posted by GitBox <gi...@apache.org>.
masahi commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r988620495
##########
python/tvm/meta_schedule/relay_integration.py:
##########
@@ -15,28 +15,101 @@
# specific language governing permissions and limitations
# under the License.
"""MetaSchedule-Relay integration"""
-from typing import Any, Dict, List, Optional
+from contextlib import contextmanager
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
import numpy as np # type: ignore
from tvm import nd
from tvm._ffi import get_global_func
from tvm.ir import IRModule, transform
from tvm.runtime import NDArray
from tvm.target import Target
+from .builder import Builder
+from .cost_model import CostModel
+from .database import Database
from .extracted_task import ExtractedTask
-from .utils import autotvm_silencer
+from .logging import get_loggers_from_work_dir
+from .measure_callback import MeasureCallback
+from .profiler import Profiler
+from .runner import Runner
+from .search_strategy import SearchStrategy
+from .space_generator import SpaceGenerator
+from .task_scheduler import TaskScheduler
+from .tune import tune_tasks
+from .tune_context import TuneContext
+from .utils import fork_seed
+
+if TYPE_CHECKING:
+ from tvm import relay
+
+_extract_task = get_global_func( # pylint: disable=invalid-name
+ "relay.backend.MetaScheduleExtractTask",
+ allow_missing=True,
+)
+
+
+@contextmanager
+def _autotvm_silencer():
+ """A context manager that silences autotvm warnings."""
+ from tvm import autotvm # pylint: disable=import-outside-toplevel
+
+ silent = autotvm.GLOBAL_SCOPE.silent
+ autotvm.GLOBAL_SCOPE.silent = True
+ try:
+ yield
+ finally:
+ autotvm.GLOBAL_SCOPE.silent = silent
+
+
+def _normalize_params(
+ mod: IRModule,
+ target: Union[Target, str],
+ params: Optional[Dict[str, NDArray]],
+ pass_config: Mapping[str, Any],
+ executor: Optional["relay.backend.Executor"],
+) -> Tuple[IRModule, Target, Dict[str, NDArray], Dict[str, Any], "relay.backend.Executor"]:
+ from tvm import relay # pylint: disable=import-outside-toplevel
+
+ if isinstance(mod, relay.Function):
+ mod = IRModule.from_expr(mod)
+ if not isinstance(target, Target):
+ target = Target(target)
+ if params is None:
+ params = {}
+ relay_params = {}
+ for name, param in params.items():
+ if isinstance(param, np.ndarray):
+ param = nd.array(param)
+ relay_params[name] = param
+ if executor is None:
+ executor = relay.backend.Executor("graph")
+ pass_config = dict(pass_config)
+ pass_config.setdefault(
+ "relay.FuseOps.link_params",
+ executor.attrs.get("link_params", False),
+ )
Review Comment:
Instead of modifying `pass_config`, we can use `mod.with_attr("executor", executor)` similarly to https://github.com/apache/tvm/blob/d4bf9ecf5524d265916ac7b860b0027f5eee5c49/src/relay/backend/build_module.cc#L414-L415. This lets us align `link-params` attribute lookup in task extraction and `relay.build`.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #12895: [MetaSchedule] Tuning API cleanup & ergonomics
Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r990232401
##########
tests/python/contrib/test_hexagon/test_meta_schedule.py:
##########
@@ -303,23 +309,20 @@ def test_vrmpy_dense_auto_tensorize(hexagon_launcher):
if True:
with tempfile.TemporaryDirectory() as work_dir:
- config = ms.TuneConfig(
- strategy="replay_trace",
+ target = get_hexagon_target("v68")
+ database = ms.tir_integration.tune_tir(
Review Comment:
this should work:
```
target = get_hexagon_target("v68")
database = ms.tir_integration.tune_tir(
mod=workload,
target=target,
max_trials_global=8,
num_trials_per_iter=8,
max_trials_per_task=8,
work_dir=work_dir,
space=ms.space_generator.PostOrderApply(
f_block_filter=None,
sch_rules=sch_rules,
postprocs=postprocs,
mutator_probs={},
),
builder=get_hexagon_local_builder(),
runner=get_hexagon_rpc_runner(hexagon_launcher, number=10),
)
sch = ms.tir_integration.compile_tir(database, workload, target)
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org