wehub-resource-sync
348 lines
13 KiB
Python
348 lines
13 KiB
Python
# 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.
|
|
"""Pre-defined pipelines.
|
|
|
|
oRelax enables flexible pipeline optimizations before min build.
|
|
This namespace offers a pre-defined collection that can be used
|
|
as it is or serves as a basis to do further composition.
|
|
"""
|
|
|
|
# pylint: disable=unused-argument
|
|
|
|
import tvm
|
|
from tvm.s_tir import meta_schedule as ms
|
|
|
|
from . import backend, transform
|
|
from .backend.utils import BackendDispatcher
|
|
|
|
|
|
def zero_pipeline(*, enable_warning: bool = False):
|
|
"""Wrapper function that returns the zero pipeline.
|
|
|
|
Parameters
|
|
----------
|
|
enable_warning : bool
|
|
A boolean value indicating if to print warnings
|
|
* in LegalizeOps pass, for CallNode whose op's legalization function is
|
|
not registered,
|
|
* in MetaScheduleApplyDatabase pass, for TIR functions now showing up in
|
|
the database. By default we don't print warning.
|
|
"""
|
|
|
|
@tvm.transform.module_pass(opt_level=0)
|
|
def f_zero_pipeline(mod: tvm.ir.IRModule, ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
|
"""Pipeline that applies pre-tuned logs.
|
|
|
|
Parameters
|
|
----------
|
|
mod : tvm.ir.IRModule
|
|
Input IRModule.
|
|
|
|
ctx : tvm.transform.PassContext
|
|
The pass context
|
|
|
|
Returns
|
|
-------
|
|
mod: tvm.ir.IRModule
|
|
The result transformed module.
|
|
"""
|
|
seq = tvm.transform.Sequential(
|
|
[
|
|
transform.LegalizeOps(enable_warning=enable_warning),
|
|
transform.AnnotateTIROpPattern(),
|
|
transform.FoldConstant(),
|
|
transform.FuseOps(),
|
|
transform.FuseTIR(),
|
|
]
|
|
)
|
|
mod = seq(mod)
|
|
if ms.Database.current():
|
|
mod = transform.MetaScheduleApplyDatabase(enable_warning=enable_warning)(mod)
|
|
return mod
|
|
|
|
return f_zero_pipeline
|
|
|
|
|
|
def default_build_pipeline():
|
|
"""The default compilation pipeline used in tvm.compile"""
|
|
|
|
@tvm.transform.module_pass(opt_level=0)
|
|
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
|
seq = tvm.transform.Sequential(
|
|
[
|
|
backend.DispatchSampling(),
|
|
backend.DispatchSortScan(),
|
|
transform.LegalizeOps(),
|
|
transform.RewriteDataflowReshape(),
|
|
transform.ToNonDataflow(),
|
|
transform.RemovePurityChecking(),
|
|
transform.CallTIRRewrite(),
|
|
transform.StaticPlanBlockMemory(),
|
|
transform.RewriteCUDAGraph(),
|
|
transform.LowerAllocTensor(),
|
|
transform.KillAfterLastUse(),
|
|
transform.LowerRuntimeBuiltin(),
|
|
transform.ComputePrimValue(),
|
|
transform.VMShapeLower(),
|
|
transform.AttachGlobalSymbol(),
|
|
],
|
|
)
|
|
mod = seq(mod)
|
|
return mod
|
|
|
|
return _pipeline
|
|
|
|
|
|
def static_shape_tuning_pipeline(
|
|
total_trials: int,
|
|
target: str | tvm.target.Target,
|
|
work_dir: str = "tuning_logs",
|
|
cpu_weight_prepack: bool = False,
|
|
max_trials_per_task: int | None = None,
|
|
):
|
|
"""Tune the static shape model and store the log to database.
|
|
|
|
Parameters
|
|
----------
|
|
total_trials : int
|
|
Total number of trials to run.
|
|
|
|
target : Union[str, tvm.target.Target]
|
|
The target device to tune the model.
|
|
|
|
work_dir : str
|
|
The directory to store the tuning logs.
|
|
|
|
cpu_weight_prepack : bool
|
|
Whether to enable the cpu weight prepack feature.
|
|
|
|
max_trials_per_task : Optional[int]
|
|
The maximum number of trials to run per task.
|
|
If not specified, it defaults to the value of `total_trials`, and this
|
|
may lead to undersubscribed tuning, potentially skipping some tasks
|
|
entirely. Explicitly setting both parameters avoids this issue and
|
|
provides deterministic resource allocation across all tasks.
|
|
For optimal tuning, set `total_trials` to at least
|
|
`max_trials_per_task * number_of_tuning_tasks` to ensure
|
|
each task receives adequate tuning resources in one iteration.
|
|
|
|
Note
|
|
----
|
|
`cpu_weight_prepack` is expected to be `True` when running on CPU for
|
|
better performance. However, it requires an explicit layout transformation
|
|
step by calling the corresponding vm function, which changes the interface
|
|
of deployment. So we disable it by default. Here is an example to enable it:
|
|
|
|
.. code-block:: python
|
|
|
|
mod = relax.pipeline.static_shape_tuning_pipeline(
|
|
total_trials=1000,
|
|
target={"kind": "llvm", "num-cores": 16},
|
|
work_dir="tuning_logs",
|
|
cpu_weight_prepack=True,
|
|
max_trials_per_task=64,
|
|
)(mod)
|
|
|
|
ex = tvm.compile(mod, target=target)
|
|
vm = relax.VirtualMachine(ex, device=tvm.cpu())
|
|
|
|
# Transform the params using the vm function
|
|
# the name should be f"{func_name}_transform_params"
|
|
params = vm["main_transform_params"](params["main"])
|
|
|
|
input_data = tvm.runtime.tensor(np.random.randn(1, 3, 224, 224).astype("float32"))
|
|
out = vm["main"](input_data, *params).numpy()
|
|
"""
|
|
|
|
@tvm.transform.module_pass(opt_level=0)
|
|
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
|
|
if cpu_weight_prepack:
|
|
pre_tuning_layout_rewrite = [transform.AttachAttrLayoutFreeBuffers()]
|
|
post_tuning_layout_rewrite = [
|
|
transform.SplitLayoutRewritePreproc(),
|
|
transform.LiftTransformParams(),
|
|
transform.FoldConstant(),
|
|
]
|
|
else:
|
|
pre_tuning_layout_rewrite = []
|
|
post_tuning_layout_rewrite = []
|
|
|
|
with tvm.target.Target(target):
|
|
mod = tvm.transform.Sequential(
|
|
[
|
|
transform.DecomposeOpsForInference(),
|
|
transform.CanonicalizeBindings(),
|
|
zero_pipeline(),
|
|
*pre_tuning_layout_rewrite,
|
|
# Skip tuning if total_trials is 0
|
|
(
|
|
transform.MetaScheduleTuneIRMod(
|
|
params={},
|
|
work_dir=work_dir,
|
|
max_trials_global=total_trials,
|
|
max_trials_per_task=max_trials_per_task,
|
|
)
|
|
if total_trials > 0
|
|
else tvm.transform.Sequential([])
|
|
),
|
|
transform.MetaScheduleApplyDatabase(work_dir),
|
|
*post_tuning_layout_rewrite,
|
|
]
|
|
)(mod)
|
|
|
|
return mod
|
|
|
|
return _pipeline
|
|
|
|
|
|
# global map of pre-built pipelines
|
|
PIPELINE_MAP = {
|
|
"zero": zero_pipeline,
|
|
"default": default_build_pipeline,
|
|
"default_build": default_build_pipeline,
|
|
"static_shape_tuning": static_shape_tuning_pipeline,
|
|
}
|
|
|
|
|
|
def get_pipeline(name: str = "zero", **kwargs) -> tvm.transform.Pass:
|
|
"""Get pre-build pipeline by name
|
|
|
|
Parameters
|
|
----------
|
|
name : Optional[str]
|
|
Name of the pipeline
|
|
kwargs : Dict[str, object]
|
|
Keyword args for configuring the pipeline.
|
|
|
|
Returns
|
|
-------
|
|
pipeline: tvm.transform.Pass
|
|
The transformation pipeline.
|
|
"""
|
|
|
|
if name not in PIPELINE_MAP:
|
|
raise ValueError(
|
|
f"Unknown pre-built pipeline {name},candidates are {list(PIPELINE_MAP.keys())}"
|
|
)
|
|
return PIPELINE_MAP[name](**kwargs)
|
|
|
|
|
|
def register_pipeline(name: str):
|
|
"""Register a new pipeline
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
Name of the pipeline
|
|
"""
|
|
if name in PIPELINE_MAP:
|
|
raise ValueError(f"Pipeline {name} has already been registered")
|
|
|
|
def _register(func):
|
|
PIPELINE_MAP[name] = func
|
|
return func
|
|
|
|
return _register
|
|
|
|
|
|
def library_dispatch_passes(target: tvm.target.Target):
|
|
"""Get the default library dispatch passes for the given target."""
|
|
if target.kind.name == "cuda":
|
|
return backend.cuda.library_dispatch_passes(target)
|
|
if target.kind.name == "rocm":
|
|
return backend.rocm.library_dispatch_passes(target)
|
|
if target.kind.name == "metal":
|
|
return backend.gpu_generic.library_dispatch_passes(target)
|
|
if target.kind.name == "llvm":
|
|
return backend.cpu_generic.library_dispatch_passes(target)
|
|
if target.kind.name == "opencl" and "adreno" in target.keys:
|
|
return backend.adreno.library_dispatch_passes(target)
|
|
if BackendDispatcher.is_gpu_target(target):
|
|
return backend.gpu_generic.library_dispatch_passes(target)
|
|
raise ValueError(f"Target {target} is not yet supported by library dispatch passes.")
|
|
|
|
|
|
def legalize_passes(target: tvm.target.Target):
|
|
"""Get the default legalization passes for the given target."""
|
|
if target.kind.name == "cuda":
|
|
return backend.cuda.legalize_passes(target)
|
|
if target.kind.name == "rocm":
|
|
return backend.rocm.legalize_passes(target)
|
|
if target.kind.name == "metal":
|
|
return backend.gpu_generic.legalize_passes(target)
|
|
if target.kind.name == "llvm":
|
|
return backend.cpu_generic.legalize_passes(target)
|
|
if target.kind.name == "opencl" and "adreno" in target.keys:
|
|
return backend.adreno.legalize_passes(target)
|
|
if BackendDispatcher.is_gpu_target(target):
|
|
return backend.gpu_generic.legalize_passes(target)
|
|
raise ValueError(f"Target {target} is not yet supported by legalize passes.")
|
|
|
|
|
|
def dataflow_lower_passes(target: tvm.target.Target):
|
|
"""Get the default legalization passes for the given target."""
|
|
if target.kind.name == "cuda":
|
|
return backend.cuda.dataflow_lower_passes(target)
|
|
if target.kind.name == "rocm":
|
|
return backend.rocm.dataflow_lower_passes(target)
|
|
if target.kind.name == "metal":
|
|
return backend.gpu_generic.dataflow_lower_passes(target)
|
|
if target.kind.name == "llvm":
|
|
return backend.cpu_generic.dataflow_lower_passes(target)
|
|
if target.kind.name == "opencl" and "adreno" in target.keys:
|
|
return backend.adreno.dataflow_lower_passes(target)
|
|
if BackendDispatcher.is_gpu_target(target):
|
|
return backend.gpu_generic.dataflow_lower_passes(target)
|
|
raise ValueError(f"Target {target} is not yet supported by dataflow lowering passes.")
|
|
|
|
|
|
def finalize_passes(target: tvm.target.Target):
|
|
"""Get the default legalization passes for the given target."""
|
|
if target.kind.name == "cuda":
|
|
return backend.cuda.finalize_passes(target)
|
|
if target.kind.name == "rocm":
|
|
return backend.rocm.finalize_passes(target)
|
|
if target.kind.name == "metal":
|
|
return backend.gpu_generic.finalize_passes(target)
|
|
if target.kind.name == "llvm":
|
|
return backend.cpu_generic.finalize_passes(target)
|
|
if target.kind.name == "opencl" and "adreno" in target.keys:
|
|
return backend.adreno.finalize_passes(target)
|
|
if BackendDispatcher.is_gpu_target(target):
|
|
return backend.gpu_generic.finalize_passes(target)
|
|
raise ValueError(f"Target {target} is not yet supported by finalization passes.")
|
|
|
|
|
|
def get_default_pipeline(target: tvm.target.Target):
|
|
"""Get the default Relax compilation pipeline for the given target."""
|
|
if target.kind.name == "cuda":
|
|
return backend.cuda.get_default_pipeline(target)
|
|
if target.kind.name == "rocm":
|
|
return backend.rocm.get_default_pipeline(target)
|
|
if target.kind.name == "metal":
|
|
return backend.gpu_generic.get_default_pipeline(target)
|
|
if target.kind.name == "llvm":
|
|
return backend.cpu_generic.get_default_pipeline(target)
|
|
if target.kind.name == "opencl" and "adreno" in target.keys:
|
|
return backend.adreno.get_default_pipeline(target)
|
|
if BackendDispatcher.is_gpu_target(target):
|
|
return backend.gpu_generic.get_default_pipeline(target)
|
|
raise ValueError(
|
|
f"Target {target} is not yet supported by default pipeline. "
|
|
"Please lower and build the IRModule manually."
|
|
)
|