wehub-resource-sync
728 lines
22 KiB
Python
728 lines
22 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.
|
|
# 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
|