wehub-resource-sync
151 lines
5.0 KiB
Python
151 lines
5.0 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.
|
|
|
|
"""Relax script for attention module."""
|
|
|
|
import tvm
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
from tvm.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import relax as relax_builder
|
|
|
|
|
|
def get_relax_attention_module(
|
|
q_shape,
|
|
k_shape,
|
|
v_shape,
|
|
*,
|
|
dtype,
|
|
bias_shape=None,
|
|
qk_scale=None,
|
|
causal_mask=None,
|
|
window_size=None,
|
|
): # pylint: disable=too-many-arguments, too-many-locals, invalid-name
|
|
"""Get a relax module for attention."""
|
|
|
|
if qk_scale is not None:
|
|
qk_scale = T.FloatImm("float32", qk_scale)
|
|
|
|
if window_size is not None:
|
|
window_size = T.IntImm("int32", window_size)
|
|
|
|
with IRBuilder() as builder:
|
|
with relax_builder.function():
|
|
R.func_name("main")
|
|
q = R.arg("q", R.Tensor(q_shape, dtype))
|
|
k = R.arg("k", R.Tensor(k_shape, dtype))
|
|
v = R.arg("v", R.Tensor(v_shape, dtype))
|
|
bias = None
|
|
if bias_shape is not None and bias_shape != "none":
|
|
bias = R.arg("bias", R.Tensor(bias_shape, dtype))
|
|
|
|
with R.dataflow() as frame:
|
|
result = R.emit(R.nn.attention(q, k, v, bias, qk_scale, causal_mask, window_size))
|
|
R.output(result)
|
|
|
|
R.func_ret_value(frame.output_vars[0])
|
|
|
|
func = builder.get()
|
|
return tvm.IRModule({"main": func})
|
|
|
|
|
|
def get_relax_stacked_attention_module(
|
|
qkv,
|
|
b,
|
|
s,
|
|
n,
|
|
h,
|
|
h_v,
|
|
op,
|
|
bias=None,
|
|
qk_scale=None,
|
|
single_shape=False,
|
|
layout="BS3NH",
|
|
): # pylint: disable=too-many-arguments, too-many-locals, too-many-branches, invalid-name
|
|
# pylint: disable=too-many-statements
|
|
"""Get a relax module for stacked attention."""
|
|
dtype = str(qkv.dtype)
|
|
assert layout in ["BS3NH", "SBN3H"]
|
|
|
|
if qk_scale is not None:
|
|
qk_scale = T.FloatImm("float32", qk_scale)
|
|
|
|
if single_shape:
|
|
if layout == "BS3NH":
|
|
qk_shape = R.shape([b, s, n, h])
|
|
elif layout == "SBN3H":
|
|
qk_shape = R.shape([b, s, n, h])
|
|
v_shape = qk_shape
|
|
else:
|
|
if layout == "BS3NH":
|
|
qk_shape = [b, s, n, h]
|
|
v_shape = [b, s, n, h_v]
|
|
elif layout == "SBN3H":
|
|
qk_shape = [s, b, n, h]
|
|
v_shape = [s, b, n, h_v]
|
|
|
|
if layout == "BS3NH":
|
|
split_axis = 2
|
|
split_sections = [n * h, n * h * 2]
|
|
elif layout == "SBN3H":
|
|
split_axis = 3
|
|
split_sections = [h, h * 2]
|
|
|
|
with IRBuilder() as builder:
|
|
with relax_builder.function():
|
|
R.func_name("main")
|
|
qkv = R.arg("qkv", R.Tensor(qkv.shape, dtype))
|
|
if bias is not None:
|
|
bias = R.arg("bias", R.Tensor(bias.shape, dtype))
|
|
with R.dataflow() as frame:
|
|
if op == "split":
|
|
qkv_tuple = R.split(qkv, split_sections, axis=split_axis)
|
|
q = qkv_tuple[0]
|
|
k = qkv_tuple[1]
|
|
v = qkv_tuple[2]
|
|
elif op == "strided_slice":
|
|
q = R.strided_slice(qkv, [split_axis], [0], [split_sections[0]], [1])
|
|
k = R.strided_slice(
|
|
qkv, [split_axis], [split_sections[0]], [split_sections[1]], [1]
|
|
)
|
|
v = R.strided_slice(
|
|
qkv,
|
|
[split_axis],
|
|
[split_sections[1]],
|
|
[int(qkv.ty.shape[split_axis])],
|
|
[1],
|
|
)
|
|
else:
|
|
raise NotImplementedError()
|
|
if layout == "BS3NH":
|
|
q = R.reshape(q, qk_shape)
|
|
k = R.reshape(k, qk_shape)
|
|
v = R.reshape(v, v_shape)
|
|
elif layout == "SBN3H":
|
|
q = R.permute_dims(q, [1, 0, 2, 3])
|
|
k = R.permute_dims(k, [1, 0, 2, 3])
|
|
v = R.permute_dims(v, [1, 0, 2, 3])
|
|
result = R.emit(R.nn.attention(q, k, v, bias, qk_scale))
|
|
if layout == "SBN3H":
|
|
result = R.emit(R.permute_dims(result, [1, 0, 2, 3]))
|
|
R.output(result)
|
|
|
|
R.func_ret_value(frame.output_vars[0])
|
|
|
|
func = builder.get()
|
|
return tvm.IRModule({"main": func})
|