This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4fdf4ae460 [Relax] Initial setup of relax backend pipeline
4fdf4ae460 is described below
commit 4fdf4ae46040e0867e9dab0d3f1877fbe14a21f1
Author: Ruihang Lai <[email protected]>
AuthorDate: Wed Feb 5 15:05:30 2025 -0500
[Relax] Initial setup of relax backend pipeline
This PR sets up the initial relax backend pipeline, as per thread
in https://discuss.tvm.apache.org/t/relax-target-default-pipeline.
Particularly, we organized default list of passes in the default
pipeline into four stages:
* library dispatch passes,
* legalize passes,
* dataflow lower passes,
* finalize passes.
The expectation is to have the default pipeline work out of the box,
while users can still choose to customize the pipeline by importing
these lists and combining them together, with customized passes
inserted at proper locations.
In this PR, we set up the initial pipelines for CUDA and LLVM backends.
The pipelines for both backends will be enhanced and completed in the
near future, and meanwhile the default pipeline for more backends
(e.g., generic GPUs) will be added.
The existing pipeline registry system in python/tvm/relax/pipeline.py
will be gradually deprecated after the establishment of the new
relax backend pipeline.
---
python/tvm/relax/backend/__init__.py | 19 +++++
.../relax/backend/{ => cpu_generic}/__init__.py | 13 ++--
python/tvm/relax/backend/cpu_generic/pipeline.py | 76 +++++++++++++++++++
python/tvm/relax/backend/{ => cuda}/__init__.py | 13 ++--
python/tvm/relax/backend/cuda/pipeline.py | 88 ++++++++++++++++++++++
tests/python/relax/test_pipeline.py | 13 ++--
6 files changed, 204 insertions(+), 18 deletions(-)
diff --git a/python/tvm/relax/backend/__init__.py
b/python/tvm/relax/backend/__init__.py
index 6d0ca30201..6b411d356d 100644
--- a/python/tvm/relax/backend/__init__.py
+++ b/python/tvm/relax/backend/__init__.py
@@ -16,7 +16,26 @@
# under the License.
"""Relax backends"""
+from tvm.target import Target
+
from . import contrib
from .dispatch_sampling import DispatchSampling
from .dispatch_sort_scan import DispatchSortScan
from .pattern_registry import get_pattern, get_patterns_with_prefix
+
+
+def get_default_pipeline(target: Target):
+ """Get the default Relax compilation pipeline for the given target."""
+ if target.kind.name == "cuda":
+ from . import cuda # pylint: disable=import-outside-toplevel
+
+ return cuda.get_default_pipeline(target)
+ if target.kind.name == "llvm":
+ from . import cpu_generic # pylint: disable=import-outside-toplevel
+
+ return cpu_generic.get_default_pipeline(target)
+ # Todo(tvm-team): support gpu-generic
+ raise ValueError(
+ f"Target {target} is not yet supported by default pipeline. "
+ "Please lower and build the IRModule manually."
+ )
diff --git a/python/tvm/relax/backend/__init__.py
b/python/tvm/relax/backend/cpu_generic/__init__.py
similarity index 78%
copy from python/tvm/relax/backend/__init__.py
copy to python/tvm/relax/backend/cpu_generic/__init__.py
index 6d0ca30201..e1cd26686c 100644
--- a/python/tvm/relax/backend/__init__.py
+++ b/python/tvm/relax/backend/cpu_generic/__init__.py
@@ -14,9 +14,10 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Relax backends"""
-
-from . import contrib
-from .dispatch_sampling import DispatchSampling
-from .dispatch_sort_scan import DispatchSortScan
-from .pattern_registry import get_pattern, get_patterns_with_prefix
+"""The Relax CPU backend compilation pipeline and other passes."""
+from .pipeline import (
+ finalize_passes,
+ get_default_pipeline,
+ legalize_passes,
+ library_dispatch_passes,
+)
diff --git a/python/tvm/relax/backend/cpu_generic/pipeline.py
b/python/tvm/relax/backend/cpu_generic/pipeline.py
new file mode 100644
index 0000000000..74d951b817
--- /dev/null
+++ b/python/tvm/relax/backend/cpu_generic/pipeline.py
@@ -0,0 +1,76 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The Relax CPU backend compilation pipeline and other passes."""
+import tvm
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default library dispatch passes for CPU backend."""
+ return []
+
+
+def legalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default legalization passes for CPU backend."""
+ return [
+ tvm.relax.transform.LegalizeOps(),
+ tvm.relax.transform.AnnotateTIROpPattern(),
+ tvm.relax.transform.FoldConstant(),
+ tvm.relax.transform.FuseOps(),
+ tvm.relax.transform.FuseTIR(),
+ ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default dataflow lowering passes for CPU backend."""
+ return [
+ relax.transform.RewriteDataflowReshape(),
+ relax.transform.ToNonDataflow(),
+ relax.transform.RemovePurityChecking(),
+ relax.transform.CallTIRRewrite(),
+ ]
+
+
+def finalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default finalization passes for CPU backend."""
+ return [
+ relax.transform.StaticPlanBlockMemory(),
+ relax.transform.LowerAllocTensor(),
+ relax.transform.KillAfterLastUse(),
+ relax.transform.LowerRuntimeBuiltin(),
+ relax.transform.ComputePrimValue(),
+ relax.transform.VMShapeLower(),
+ relax.transform.AttachGlobalSymbol(),
+ ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+ """Return the default compilation pipeline for CPU."""
+
+ @tvm.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ with target:
+ seq = tvm.transform.Sequential(
+ library_dispatch_passes(target)
+ + legalize_passes(target)
+ + dataflow_lower_passes(target)
+ + finalize_passes(target)
+ )
+ mod = seq(mod)
+ return mod
+
+ return _pipeline
diff --git a/python/tvm/relax/backend/__init__.py
b/python/tvm/relax/backend/cuda/__init__.py
similarity index 78%
copy from python/tvm/relax/backend/__init__.py
copy to python/tvm/relax/backend/cuda/__init__.py
index 6d0ca30201..f4458f4b55 100644
--- a/python/tvm/relax/backend/__init__.py
+++ b/python/tvm/relax/backend/cuda/__init__.py
@@ -14,9 +14,10 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Relax backends"""
-
-from . import contrib
-from .dispatch_sampling import DispatchSampling
-from .dispatch_sort_scan import DispatchSortScan
-from .pattern_registry import get_pattern, get_patterns_with_prefix
+"""The Relax CUDA backend compilation pipeline and other passes."""
+from .pipeline import (
+ finalize_passes,
+ get_default_pipeline,
+ legalize_passes,
+ library_dispatch_passes,
+)
diff --git a/python/tvm/relax/backend/cuda/pipeline.py
b/python/tvm/relax/backend/cuda/pipeline.py
new file mode 100644
index 0000000000..d5c4c08561
--- /dev/null
+++ b/python/tvm/relax/backend/cuda/pipeline.py
@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The Relax CUDA backend compilation pipeline and other passes."""
+import tvm
+from tvm import dlight as dl
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default library dispatch passes for CUDA backend."""
+ return [
+ relax.backend.DispatchSampling(),
+ relax.backend.DispatchSortScan(),
+ ]
+
+
+def legalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default legalization passes for CUDA backend."""
+ return [
+ tvm.relax.transform.LegalizeOps(),
+ tvm.relax.transform.AnnotateTIROpPattern(),
+ tvm.relax.transform.FoldConstant(),
+ tvm.relax.transform.FuseOps(),
+ tvm.relax.transform.FuseTIR(),
+ dl.ApplyDefaultSchedule(
+ dl.gpu.Matmul(),
+ dl.gpu.GEMV(),
+ dl.gpu.Reduction(),
+ dl.gpu.GeneralReduction(),
+ dl.gpu.Fallback(),
+ ),
+ ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default dataflow lowering passes for CUDA backend."""
+ return [
+ relax.transform.RewriteDataflowReshape(),
+ relax.transform.ToNonDataflow(),
+ relax.transform.RemovePurityChecking(),
+ relax.transform.CallTIRRewrite(),
+ ]
+
+
+def finalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
+ """The default finalization passes for CUDA backend."""
+ return [
+ relax.transform.StaticPlanBlockMemory(),
+ relax.transform.RewriteCUDAGraph(),
+ relax.transform.LowerAllocTensor(),
+ relax.transform.KillAfterLastUse(),
+ relax.transform.LowerRuntimeBuiltin(),
+ relax.transform.ComputePrimValue(),
+ relax.transform.VMShapeLower(),
+ relax.transform.AttachGlobalSymbol(),
+ ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+ """Return the default compilation pipeline for CUDA."""
+
+ @tvm.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ with target:
+ seq = tvm.transform.Sequential(
+ library_dispatch_passes(target)
+ + legalize_passes(target)
+ + dataflow_lower_passes(target)
+ + finalize_passes(target)
+ )
+ mod = seq(mod)
+ return mod
+
+ return _pipeline
diff --git a/tests/python/relax/test_pipeline.py
b/tests/python/relax/test_pipeline.py
index 8aa41490c6..86711477b2 100644
--- a/tests/python/relax/test_pipeline.py
+++ b/tests/python/relax/test_pipeline.py
@@ -15,14 +15,17 @@
# specific language governing permissions and limitations
# under the License.
import numpy as np
+
import tvm
import tvm.testing
from tvm import relax
-from tvm.script import relax as R, tir as T
+from tvm.script import relax as R
+from tvm.script import tir as T
def test_pipeline_compile():
- pipeline = relax.get_pipeline()
+ target = tvm.target.Target("llvm", host="llvm")
+ pipeline = relax.backend.get_default_pipeline(target)
@tvm.script.ir_module
class Mod:
@@ -33,7 +36,6 @@ def test_pipeline_compile():
mod = Mod
mod = pipeline(mod)
- target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(mod, target)
x_np = np.random.rand(3, 4).astype(np.float32)
@@ -48,7 +50,8 @@ def test_pipeline_compile():
def test_pipeline_with_kv_cache():
"""A dummy pipline that simulates KV update."""
- pipeline = relax.get_pipeline()
+ target = tvm.target.Target("llvm", host="llvm")
+ pipeline = relax.backend.get_default_pipeline(target)
@tvm.script.ir_module
class Mod:
@@ -92,8 +95,6 @@ def test_pipeline_with_kv_cache():
mod = Mod
mod = pipeline(mod)
- target = tvm.target.Target("llvm", host="llvm")
-
ex = relax.build(mod, target)
num_steps = 8
