mirrors/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/backend/adreno/__init__.py
+28
View File
@@ -0,0 +1,28 @@
# isort: skip_file
# 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 Adreno backend compilation pipeline and other passes."""
from . import transform
from .pipeline import (
finalize_passes,
get_default_pipeline,
dataflow_lower_passes,
legalize_passes,
library_dispatch_passes,
)
python/tvm/relax/backend/adreno/clml.py
+727
View File
@@ -0,0 +1,727 @@
# 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=invalid-name, unused-argument, pointless-exception-statement
"""Pattern table for CLML backend"""
import tvm
from tvm import IRModule, relax, tirx
from tvm.ir.transform import PassContext, module_pass
from tvm.relax import transform
from tvm.relax.dpl.pattern import (
GlobalVarPattern,
TuplePattern,
is_const,
is_op,
is_tuple_get_item,
wildcard,
)
from tvm.relax.expr import TupleGetItem, VarBinding
from tvm.relax.expr_functor import PyExprMutator, mutator
from tvm.relax.transform import PatternCheckContext
from ..pattern_registry import register_patterns
def _dtype_str(dtype):
return str(dtype.dtype) if isinstance(dtype, tvm.ir.PrimType) else str(dtype)
@mutator
class AppendReshapeToBNRewriter(PyExprMutator):
"""
Append Reshape Operator to BatchNorm Pass Rewriter Pass
- Automatically appends a reshape operation after BatchNorm operators
- Resolves fusion issues for custom backends where BatchNorm output
might explicitly access the first elment of the Tuple
Algo:
Identifies BatchNorm operators in the computational graph
When BatchNorm's first output is accessed via TupleGetItem
Automatically inserts a reshape operation to match input shape
"""
def __init__(self, mod):
super().__init__(mod)
self.bn_vars = {}
def visit_tuple_getitem_(self, op: TupleGetItem):
tuple_value = op.tuple_value
reshape_op = tvm.ir.Op.get("relax.reshape")
if isinstance(tuple_value, relax.Var) and tuple_value in self.bn_vars:
bn_call = self.bn_vars[tuple_value]
if op.index == 0:
bn_out = relax.TupleGetItem(bn_call, 0)
input_shape = bn_call.args[0].ty.shape
return relax.Call(reshape_op, [bn_out, input_shape])
return super().visit_tuple_getitem_(op)
def visit_var_binding_(self, binding: VarBinding):
if isinstance(binding.value, relax.Call) and binding.value.op.name == "relax.nn.batch_norm":
self.bn_vars[binding.var] = binding.value
return super().visit_var_binding_(binding)
@transform.function_pass(opt_level=0, name="AppendReshapeToBN")
class AppendReshapeToBNRewriterPass:
def transform_function(
self, func: relax.Function, mod: IRModule, _ctx: tvm.transform.PassContext
) -> relax.Function:
updated_func = AppendReshapeToBNRewriter(mod).visit_expr(func)
updated_func = relax.analysis.remove_all_unused(updated_func)
return updated_func
def clml_sdk_version():
"""Utility function to get clml version.
Probes the FFI registry for the OpenCLML version registered by the
CLML backend at build time. Returns 2 when CLML is not present.
"""
# Registry: "relax.get_openclml_version" — returns the CLML SDK version
# that TVM was built against; registered unconditionally in codegen.cc.
# Grep hint: grep -rn 'relax.get_openclml_version' src/
get_version = tvm.get_global_func("relax.get_openclml_version", allow_missing=True)
if get_version is None:
return 2
return int(get_version())
def is_clml_runtime_enabled():
"""Check if the CLML graph runtime is present.
Returns
-------
ret: bool
True if present, False if not.
"""
check_enabled = tvm.get_global_func("relax.op.is_openclml_runtime_enabled", True)
if check_enabled:
return check_enabled()
return False
def _check_default(context: PatternCheckContext) -> bool:
return True
def clml_pattern_table():
"""Get the CLML pattern table."""
def _check_conv2d(context: PatternCheckContext) -> bool:
if "root" in context.annotated_expr:
root_call = context.annotated_expr["root"]
if root_call.op.name == "relax.nn.conv2d":
input_layout = root_call.attrs.data_layout
weight_layout = root_call.attrs.kernel_layout
if input_layout != "NCHW" or weight_layout != "OIHW":
return False
if root_call.op.name == "relax.nn.conv2d_transpose":
input_layout = root_call.attrs.data_layout
weight_layout = root_call.attrs.kernel_layout
if input_layout != "NCHW" or weight_layout != "OIHW":
return False
if "data" in context.annotated_expr:
input_expr = context.annotated_expr["data"]
input_dtype = _dtype_str(input_expr.ty.dtype)
if input_dtype not in ["float32", "float16"]:
return False
if "weight" in context.annotated_expr:
weight_expr = context.annotated_expr["weight"]
weight_dtype = _dtype_str(weight_expr.ty.dtype)
if weight_dtype not in ["float32", "float16"]:
return False
return True
def populate_patterns(patterns, name, op, annotations, *args):
ret = {}
for k, v in patterns.items():
ret_ann = v["annotation"].copy()
ret_ann.update(annotations)
ret[name + "." + k] = {"pattern": op(v["pattern"], *args), "annotation": ret_ann.copy()}
return ret
def conv_pattern():
"""Create a convolution pattern."""
data = wildcard()
weight = wildcard()
bias = is_const()
bn_scale = is_const()
bn_bias = is_const()
bn_mean = is_const()
bn_var = is_const()
annotations = {
"data": data,
"weight": weight,
}
patterns = {}
patterns["nn.conv2d"] = {
"pattern": is_op("relax.nn.conv2d")(data, weight),
"annotation": annotations.copy(),
}
pad_annotations = annotations.copy()
patterns["pad.nn.conv2d"] = {
"pattern": is_op("relax.nn.conv2d")(is_op("relax.nn.pad")(data), weight),
"annotation": pad_annotations,
}
patterns["nn.conv2d_transpose"] = {
"pattern": is_op("relax.nn.conv2d_transpose")(data, weight),
"annotation": annotations.copy(),
}
patterns.update(
populate_patterns(patterns, "bias", is_op("relax.add"), {"bias": bias}, bias)
)
patterns.update(
populate_patterns(
patterns,
"bn",
is_op("relax.nn.batch_norm"),
{
"bn_scale": bn_scale,
"bn_bias": bn_bias,
"bn_mean": bn_mean,
"bn_var": bn_var,
},
bn_scale,
bn_bias,
bn_mean,
bn_var,
)
)
tuple_patterns = {}
for k, v in patterns.items():
tuple_annotation = v["annotation"].copy()
tuple_patterns["tuple" + "." + k] = {
"pattern": is_tuple_get_item(v["pattern"], 0),
"annotation": tuple_annotation,
}
patterns.update(tuple_patterns)
relu_patterns = populate_patterns(patterns, "relu", is_op("relax.nn.relu"), {})
clip_patterns = populate_patterns(patterns, "clip", is_op("relax.clip"), {})
patterns.update(relu_patterns)
patterns.update(clip_patterns)
conv_patterns = []
for k, v in patterns.items():
ret_annotations = v["annotation"]
ret_annotations["root"] = v["pattern"]
conv_patterns.append(
("openclml." + (k), v["pattern"], ret_annotations.copy(), _check_conv2d)
)
return conv_patterns[::-1]
def _check_maxpool2d(context: PatternCheckContext) -> bool:
root = context.annotated_expr.get("root")
if root is None or not isinstance(root, relax.Call):
return False
if root.op.name != "relax.nn.max_pool2d":
return False
if "data" not in context.annotated_expr:
return False
data = context.annotated_expr["data"]
input_shape = data.ty.shape
if len(input_shape) != 4:
return False
if any(dim <= 0 for dim in input_shape):
return False
pool_size = root.attrs.pool_size
if len(pool_size) != 2:
return False
if any(size <= 0 for size in pool_size):
return False
strides = root.attrs.strides
if len(strides) != 2:
return False
if any(stride <= 0 for stride in strides):
return False
dilation = root.attrs.dilation
if len(dilation) != 2:
return False
if any(d <= 0 for d in dilation):
return False
padding = root.attrs.padding
if len(padding) != 4:
return False
if any(p < 0 for p in padding):
return False
return True
def maxpool_pattern():
"""Create Pool Pattern"""
data = wildcard()
annotations = {
"data": data,
}
patterns = {}
patterns["nn.max_pool2d"] = {
"pattern": is_op("relax.nn.max_pool2d")(data),
"annotation": annotations.copy(),
}
pool_patterns = []
for k, v in patterns.items():
ret_annotations = v["annotation"]
ret_annotations["root"] = v["pattern"]
pool_patterns.append(
("openclml." + (k), v["pattern"], ret_annotations.copy(), _check_maxpool2d)
)
return pool_patterns
def _check_avgpool2d(context: PatternCheckContext) -> bool:
root = context.annotated_expr.get("root")
if root is None or not isinstance(root, relax.Call):
return False
if root.op.name != "relax.nn.avg_pool2d":
return False
if "data" not in context.annotated_expr:
return False
data = context.annotated_expr["data"]
input_shape = data.ty.shape
if len(input_shape) != 4:
return False
if any(dim <= 0 for dim in input_shape):
return False
pool_size = root.attrs.pool_size
if len(pool_size) != 2:
return False
if any(size <= 0 for size in pool_size):
return False
strides = root.attrs.strides
if len(strides) != 2:
return False
if any(stride <= 0 for stride in strides):
return False
padding = root.attrs.padding
if len(padding) != 4:
return False
if any(p < 0 for p in padding):
return False
return True
def avgpool_pattern():
data = wildcard()
annotations = {
"data": data,
}
patterns = {}
patterns["nn.avg_pool2d"] = {
"pattern": is_op("relax.nn.avg_pool2d")(data),
"annotation": annotations.copy(),
}
pool_patterns = []
for k, v in patterns.items():
ret_annotations = v["annotation"]
ret_annotations["root"] = v["pattern"]
pool_patterns.append(
("openclml." + (k), v["pattern"], ret_annotations.copy(), _check_avgpool2d)
)
return pool_patterns
def _check_global_avgpool(context: PatternCheckContext) -> bool:
root = context.annotated_expr.get("root")
if root is None or not isinstance(root, relax.Call):
return False
if root.op.name != "relax.mean":
return False
if "data" not in context.annotated_expr:
return False
data = context.annotated_expr["data"]
input_shape = data.ty.shape
if len(input_shape) != 4:
return False
if input_shape[1] <= 0 or input_shape[2] <= 0 or input_shape[3] <= 0:
return False
if not hasattr(root.attrs, "axis"):
return False
axis = root.attrs.axis
if not (len(axis) == 2 and axis[0] == 2 and axis[1] == 3):
return False
return True
def global_avgpool_pattern():
"""Create Pool Pattern"""
data = wildcard()
pattern = is_op("relax.mean")(data).has_attr({"axis": [2, 3]})
annotations = {
"data": data,
"root": pattern,
}
return [
("openclml.nn.global_avg_pool2d", pattern, annotations, _check_global_avgpool),
]
def _check_reshape(context: PatternCheckContext) -> bool:
root = context.annotated_expr.get("root")
if root is None or not isinstance(root, relax.Call):
return False
if root.op.name != "relax.reshape":
return False
shape_arg = root.args[1]
if not isinstance(shape_arg, relax.Expr):
return False
return True
def reshape_pattern():
"""Create Reshape Pattern"""
pattern = is_op("relax.reshape")(wildcard(), wildcard())
annotations = {
"root": pattern,
}
return [("openclml.reshape", pattern, annotations, _check_reshape)]
def _check_batchnorm(context: PatternCheckContext) -> bool:
root = context.annotated_expr.get("root")
if root is None or not isinstance(root, relax.Call):
return False
if root.op.name != "relax.reshape":
return False
required_params = ["moving_var", "gamma", "moving_mean", "beta"]
for param in required_params:
if param not in context.annotated_expr:
return False
params = {
"moving_var": context.annotated_expr["moving_var"],
"gamma": context.annotated_expr["gamma"],
"moving_mean": context.annotated_expr["moving_mean"],
"beta": context.annotated_expr["beta"],
}
for param in params.values():
if not isinstance(param, relax.expr.Constant):
return False
base_shape = None
for param in params.values():
shape = param.ty.shape
dtype = _dtype_str(param.ty.dtype)
if dtype not in {"float32"}:
return False
# Initialize base_shape if not set
if base_shape is None:
base_shape = shape
continue
# All parameters should have same shape
if len(shape) != len(base_shape):
return False
if any(s1 != s2 for s1, s2 in zip(shape, base_shape)):
return False
return True
def batch_norm_pattern():
"""Create a batch norm pattern."""
data = wildcard()
bn_scale = is_const()
bn_bias = is_const()
bn_mean = is_const()
bn_var = is_const()
pattern = is_op("relax.nn.batch_norm")(data, bn_scale, bn_bias, bn_mean, bn_var)
pattern = is_tuple_get_item(pattern, 0)
pattern = is_op("relax.reshape")(pattern, wildcard())
annotations = {
"gamma": bn_scale,
"beta": bn_bias,
"moving_mean": bn_mean,
"moving_var": bn_var,
"root": pattern,
}
return [
("openclml.nn.batch_norm", pattern, annotations, _check_batchnorm),
]
def _check_binary_op(context: PatternCheckContext) -> bool:
def _check_arg(input_expr):
input_dtype = _dtype_str(input_expr.ty.dtype)
input_shape = input_expr.ty.shape
if len(input_shape) == 0:
return False
# Avoid any operators with dtype Int64
if input_dtype == "int64":
return False
# No support for batch> 1
if input_shape[0] > 1:
return False
return True
def compare_shapes(lhs_shape, rhs_shape):
if len(lhs_shape) != len(rhs_shape):
return False
for lhs_dim, rhs_dim in zip(lhs_shape, rhs_shape):
if lhs_dim != rhs_dim:
return False
return True
lhs_shape = None
rhs_shape = None
if "lhs" in context.annotated_expr:
lhs = context.annotated_expr["lhs"]
lhs_shape = lhs.ty.shape
if not _check_arg(lhs):
return False
if "rhs" in context.annotated_expr:
rhs = context.annotated_expr["rhs"]
rhs_shape = rhs.ty.shape
if not _check_arg(rhs):
return False
# Checking for BinaryOps ( False for unaryOp )
if (
"lhs" in context.annotated_expr
and "rhs" in context.annotated_expr
and not compare_shapes(lhs_shape, rhs_shape)
):
return False
return True
def binary_op_pattern():
"""Create a binary op pattern."""
def make_pattern(op):
lhs = wildcard()
rhs = wildcard()
pattern = is_op(op)(lhs, rhs)
annotations = {"lhs": lhs, "rhs": rhs}
return ("openclml." + op, pattern, annotations, _check_binary_op)
binary_ops = [
"relax.add",
"relax.subtract",
"relax.multiply",
"relax.divide",
"relax.maximum",
"relax.minimum",
]
return [make_pattern(op) for op in binary_ops]
def unary_op_pattern():
"""Create a unary op pattern."""
def make_pattern(op):
lhs = wildcard()
pattern = is_op(op)(lhs)
annotations = {"lhs": lhs}
return ("openclml." + op, pattern, annotations, _check_binary_op)
unary_ops = [
"relax.nn.softmax",
"relax.nn.relu",
"relax.clip",
]
return [make_pattern(op) for op in unary_ops]
return [
*conv_pattern(),
*batch_norm_pattern(),
*binary_op_pattern(),
*unary_op_pattern(),
*maxpool_pattern(),
*avgpool_pattern(),
*global_avgpool_pattern(),
*reshape_pattern(),
]
clml_patterns = clml_pattern_table()
register_patterns(clml_patterns)
@module_pass(opt_level=0, name="OpenCLMLOffLoad")
class OpenCLMLOffLoad:
"""The pass sequence used for CLML offload"""
def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule:
"""The transform"""
clml_layouts = {
"relax.nn.conv2d": ["NCHW", "OIHW"],
"relax.nn.conv2d_transpose": ["NCHW", "OIHW"],
}
seq = tvm.transform.Sequential(
[
transform.ConvertLayout(clml_layouts),
transform.Normalize(),
transform.FoldBatchnormToConv2D(),
AppendReshapeToBNRewriterPass(),
transform.FoldConstant(),
transform.FuseOpsByPattern(clml_pattern_table()),
transform.MergeCompositeFunctions(),
transform.RunCodegen(),
],
)
mod = seq(mod)
return mod
def _check_dequantize_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
_input = ctx.annotated_expr["lhs"]
root = ctx.annotated_expr["root"]
wdq = ctx.annotated_expr["w_decoded"]
w_pack = ctx.annotated_expr["w_encoded"]
if _dtype_str(ctx.annotated_expr["lhs"].ty.dtype) != "float16":
return False
if not isinstance(wdq, relax.Call):
return False
g_var = wdq.args[0]
if not (isinstance(g_var, relax.GlobalVar) and "dequantize" in g_var.name_hint):
return False
if not (
(len(root.ty.shape) == 3)
and isinstance(root.ty.shape[0], tirx.IntImm)
and (_dtype_str(root.ty.dtype) == "float16")
and (root.ty.shape[0] == 1)
):
return False
if not (
(len(wdq.ty.shape) == 2)
and (w_pack.ty.shape[-1] == root.ty.shape[-1])
and (wdq.ty.shape[-2] == _input.ty.shape[-1])
):
return False
return True
def dequantize_matmul_patterns():
"""Returns a list of supported decode -> matmul patterns."""
def _dequantize_matmul_pattern(name):
scales = wildcard()
x = wildcard()
w_packed = wildcard()
w_decoded = is_op("relax.call_tir")(
GlobalVarPattern(),
TuplePattern([w_packed, scales]),
)
matmul = is_op("relax.matmul")(x, w_decoded)
annotations = {
"root": matmul,
"lhs": x,
"w_encoded": w_packed,
"w_decoded": w_decoded,
"scales": scales,
}
return name, matmul, annotations, _check_dequantize_matmul
return [
_dequantize_matmul_pattern("openclml.dequant_matmul"),
]
clml_llm_patterns = [
*dequantize_matmul_patterns(),
]
register_patterns(clml_llm_patterns)
@tvm.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
class OpenCLMLOffLoadForLLM:
"""A compiler pass that partition the graph with dequant Matmul to CLML backend offload."""
def __init__(self, target: tvm.target.Target) -> None:
"""Initializer.
Parameters
----------
target : tvm.target.Target
Target device.
"""
self.target = target
def transform_module(
self,
mod: IRModule,
_ctx: tvm.transform.PassContext,
) -> IRModule:
"""Apply required passed to transform"""
if "adreno" in self.target.keys and (clml_sdk_version() >= 5):
mod = tvm.transform.Sequential(
[
transform.Normalize(),
transform.FuseOpsByPattern(clml_llm_patterns, annotate_codegen=True),
transform.RunCodegen(),
]
)(mod)
return mod
python/tvm/relax/backend/adreno/pipeline.py
+138
View File
@@ -0,0 +1,138 @@
# 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 Adreno GPU backend compilation pipeline and other passes."""
import tvm
from tvm import relax
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
"""The default library dispatch passes for Adreno GPU backend."""
if "clml" in target.keys:
return [
relax.backend.adreno.clml.OpenCLMLOffLoadForLLM(target),
relax.backend.adreno.clml.OpenCLMLOffLoad(),
]
else:
return []
def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
"""The default legalization passes for Adreno GPU backend."""
desired_layouts = {"relax.nn.conv2d": ["NCHW4c", "OIHW4o", "NCHW4c"]}
skip_ops = [
"relax.nn.conv2d",
"relax.nn.max_pool2d",
"relax.nn.adaptive_avg_pool2d",
]
pass_list = []
pass_list.extend(
[
tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)),
relax.transform.DecomposeOpsForInference(),
]
)
if "texture" in target.keys:
pass_list.extend(
[
relax.transform.ConvertLayout(desired_layouts),
relax.transform.Normalize(),
relax.transform.FoldConstant(),
relax.transform.LegalizeOps(skip_ops=skip_ops),
relax.transform.AnnotateTIROpPattern(),
relax.backend.adreno.transform.AnnotateCustomMemoryScope(target),
]
)
pass_list.extend([tvm.relax.transform.LegalizeOps()])
if "texture" in target.keys:
pass_list.extend(
[
relax.transform.LegalizeOps(
{"relax.nn.conv2d": legalize_adreno.conv2d_NCHWc_OIHWo},
)
]
)
pass_list.extend(
[
relax.transform.AnnotateTIROpPattern(),
relax.transform.FoldConstant(),
relax.transform.FuseOps(),
relax.transform.FuseTIR(),
relax.transform.DeadCodeElimination(),
]
)
if "texture" in target.keys:
pass_list.extend(
[
relax.backend.adreno.transform.FoldVDeviceScopeChange(),
relax.transform.DeadCodeElimination(),
relax.transform.SpecializePrimFuncBasedOnCallSite(),
]
)
from tvm.s_tir import dlight as dl # pylint: disable=import-outside-toplevel
pass_list.extend([relax.transform.Normalize()])
pass_list.extend(
[
dl.ApplyDefaultSchedule(
dl.adreno.Conv2d(),
dl.adreno.LayoutTransform(),
dl.adreno.Pool2D(),
)
]
)
pass_list.extend(
[
dl.ApplyDefaultSchedule(
dl.gpu.Reduction(),
dl.gpu.GeneralReduction(),
dl.gpu.Fallback(),
)
]
)
return pass_list
def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
"""The default dataflow lowering passes for Adreno GPU backend."""
return relax.backend.gpu_generic.dataflow_lower_passes(target)
def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
"""The default finalization passes for Adreno GPU backend."""
return relax.backend.gpu_generic.finalize_passes(target)
def get_default_pipeline(target: tvm.target.Target):
"""Return the default compilation pipeline for Adreno GPU."""
@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
python/tvm/relax/backend/adreno/transform/__init__.py
+23
View File
@@ -0,0 +1,23 @@
# isort: skip_file
# 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.
"""Adreno Relax transformations."""
from .transform import (
AnnotateCustomMemoryScope,
FoldVDeviceScopeChange,
)
python/tvm/relax/backend/adreno/transform/_ffi_api.py
+20
View File
@@ -0,0 +1,20 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for Adreno transform"""
import tvm_ffi
tvm_ffi.init_ffi_api("relax.backend.adreno.transform", __name__)
python/tvm/relax/backend/adreno/transform/transform.py
+49
View File
@@ -0,0 +1,49 @@
# 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=invalid-name
"""Adreno Relax transformation passes."""
import tvm.ir
from tvm.target import Target
from . import _ffi_api
def AnnotateCustomMemoryScope(target: Target | None = None) -> tvm.ir.transform.Pass:
"""Allocate the memory scope information. This is Adreno specific pass to annotate
The memory scope information and realize the same with RealizeVDevice pass followed by
updating the Prim Function var_buffer mapping using SpecializePrimFuncBasedOnCallSite.
Returns
-------
ret: tvm.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.AnnotateCustomMemoryScope(target) # type: ignore
def FoldVDeviceScopeChange() -> tvm.ir.transform.Pass:
"""This pass is a texture specific pass that can optimize unnecessary to_device copies.
Like texture_scope -> ToVDevice -> global scope. In this case the producer can directly
store into global scope avoiding unnecessary device copy.
Returns
-------
ret: tvm.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.FoldVDeviceScopeChange() # type: ignore