wehub-resource-sync
310 lines
13 KiB
Python
310 lines
13 KiB
Python
"""
|
|
Implementation for Phi architecture.
|
|
"""
|
|
|
|
from tvm import relax, tirx
|
|
from tvm.relax.frontend import nn
|
|
from tvm.relax.frontend.nn import Module, Tensor, op
|
|
from tvm.script import tirx as T
|
|
|
|
from mlc_llm.model.vision import CLIPVisionModel
|
|
from mlc_llm.support.config import ConfigBase
|
|
|
|
|
|
# mypy: disable-error-code="attr-defined"
|
|
class ImageProjection(Module):
|
|
def __init__(self, config: ConfigBase):
|
|
super().__init__()
|
|
self.linear_1 = nn.Linear(
|
|
config.vision_config.hidden_size * 4, config.hidden_size, bias=True
|
|
)
|
|
self.act = nn.GELU()
|
|
self.linear_2 = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
|
|
|
|
def forward(self, image_features: Tensor) -> Tensor:
|
|
shape_1 = tirx.Var("shape_1", "int64")
|
|
image_features = op.wrap_nested(
|
|
relax.BlockBuilder()
|
|
.current()
|
|
.match_cast(
|
|
image_features._expr,
|
|
relax.TensorType([shape_1, image_features.shape[1]], image_features.dtype),
|
|
),
|
|
"image_features",
|
|
)
|
|
|
|
hidden_states = self.linear_1(image_features)
|
|
|
|
shape_2 = tirx.Var("shape_2", "int64")
|
|
hidden_states = op.wrap_nested(
|
|
relax.BlockBuilder()
|
|
.current()
|
|
.match_cast(
|
|
hidden_states._expr,
|
|
relax.TensorType([shape_2, hidden_states.shape[1]], hidden_states.dtype),
|
|
),
|
|
"hidden_states",
|
|
)
|
|
|
|
hidden_states = self.act(hidden_states)
|
|
hidden_states = self.linear_2(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class Phi3ImageEmbedding(Module):
|
|
def __init__(self, config: ConfigBase):
|
|
super().__init__()
|
|
|
|
self.img_processor = CLIPVisionModel(config.vision_config)
|
|
self.image_dim_out = config.img_processor["image_dim_out"]
|
|
|
|
self.glb_GN = nn.Parameter((1, 1, self.image_dim_out * 4))
|
|
self.sub_GN = nn.Parameter((1, 1, 1, self.image_dim_out * 4))
|
|
|
|
self.img_projection = ImageProjection(config)
|
|
self.image_size = config.vision_config.image_size
|
|
|
|
def apply_schedule(self, sch, block, bdx=32, tile=[32, 32]):
|
|
loop_x, loop_y = sch.get_loops(block)[-2:]
|
|
xo, xi = sch.split(loop_x, factors=[tile[0], None])
|
|
yo, yi = sch.split(loop_y, factors=[tile[1], None])
|
|
sch.reorder(xo, yo, xi, yi)
|
|
t = sch.fuse(xo, yo)
|
|
ty, tx = sch.split(t, factors=[None, bdx])
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
|
|
def dyn_repeat_4d_tensor(self, input_tensor, r0, r1, r2, r3) -> Tensor:
|
|
assert 4 == input_tensor.ndim, "input_tensor should be 4D data tensor"
|
|
|
|
def create_dyn_repeat_func(dtype):
|
|
@T.prim_func(s_tir=True)
|
|
def dyn_repeat_4d_tensor_func( # pylint disable=too-many-locals
|
|
input_tensor: T.handle,
|
|
output: T.handle,
|
|
ch0: T.int64(),
|
|
ch1: T.int64(),
|
|
ch2: T.int64(),
|
|
ch3: T.int64(),
|
|
):
|
|
T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
|
|
n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64()
|
|
input_tensor_buf = T.match_buffer(input_tensor, (n, c, h, w), dtype=dtype)
|
|
out_buf = T.match_buffer(output, (n * ch0, c * ch1, h * ch2, w * ch3), dtype=dtype)
|
|
|
|
for n_idx in T.thread_binding(n * ch0, thread="blockIdx.x"):
|
|
for c_idx in T.thread_binding(c * ch1, thread="blockIdx.y"):
|
|
for h_idx, w_idx in T.grid(h * ch2, w * ch3):
|
|
with T.sblock("dyn_repeat_4d_tensor"):
|
|
T.reads(input_tensor_buf[n_idx, c_idx, h_idx, w_idx])
|
|
T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
|
|
out_buf[n_idx, c_idx, h_idx, w_idx] = input_tensor_buf[
|
|
n_idx % n, c_idx % c, h_idx % h, w_idx % w
|
|
]
|
|
|
|
return dyn_repeat_4d_tensor_func
|
|
|
|
n, c, h, w = input_tensor.shape
|
|
out = op.tensor_ir_op(
|
|
create_dyn_repeat_func(input_tensor.dtype),
|
|
"dyn_repeat_4d_tensor",
|
|
[input_tensor, r0, r1, r2, r3],
|
|
[Tensor.placeholder([n * r0, c * r1, h * r2, w * r3], input_tensor.dtype)],
|
|
)
|
|
return out
|
|
|
|
def dyn_concate_dim_2(self, input_1, input_2) -> Tensor:
|
|
def create_dyn_concate_func(dtype):
|
|
@T.prim_func(s_tir=True)
|
|
def dyn_concate_dim_2_func(input_1: T.handle, input_2: T.handle, output: T.handle):
|
|
T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
|
|
n, c, h1, h2, w = T.int64(), T.int64(), T.int64(), T.int64(), T.int64()
|
|
input_1_buf = T.match_buffer(input_1, (n, c, h1, w), dtype=dtype)
|
|
input_2_buf = T.match_buffer(input_2, (n, c, h2, w), dtype=dtype)
|
|
out_buf = T.match_buffer(output, (n, c, h1 + h2, w), dtype=dtype)
|
|
|
|
for n_idx in T.thread_binding(n, thread="blockIdx.x"):
|
|
for c_idx in T.thread_binding(c, thread="blockIdx.y"):
|
|
for h_idx, w_idx in T.grid(h1 + h2, w):
|
|
with T.sblock("dyn_concate_dim_2"):
|
|
T.reads(input_1_buf[n_idx, c_idx, h_idx, w_idx])
|
|
T.writes(out_buf[n_idx, c_idx, h_idx, w_idx])
|
|
if h_idx < h1:
|
|
out_buf[n_idx, c_idx, h_idx, w_idx] = input_1_buf[
|
|
n_idx, c_idx, h_idx, w_idx
|
|
]
|
|
else:
|
|
out_buf[n_idx, c_idx, h_idx, w_idx] = input_2_buf[
|
|
n_idx, c_idx, h_idx - h1, w_idx
|
|
]
|
|
|
|
return dyn_concate_dim_2_func
|
|
|
|
n1, c1, h1, w1 = input_1.shape
|
|
n2, c2, h2, w2 = input_2.shape
|
|
assert n1 == n2 and c1 == c2 and w1 == w2
|
|
|
|
out = op.tensor_ir_op(
|
|
create_dyn_concate_func(input_1.dtype),
|
|
"dyn_concate_dim_2",
|
|
[input_1, input_2],
|
|
[Tensor.placeholder([n1, c1, h1 + h2, w1], input_1.dtype)],
|
|
)
|
|
return out
|
|
|
|
def dyn_concate_dim_1(self, input_1, input_2) -> Tensor:
|
|
def create_dyn_concate_func(dtype):
|
|
@T.prim_func(s_tir=True)
|
|
def dyn_concate_dim_1_func(input_1: T.handle, input_2: T.handle, output: T.handle):
|
|
T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1})
|
|
c, h1, h2, w = T.int64(), T.int64(), T.int64(), T.int64()
|
|
input_1_buf = T.match_buffer(input_1, (c, h1, w), dtype=dtype)
|
|
input_2_buf = T.match_buffer(input_2, (c, h2, w), dtype=dtype)
|
|
out_buf = T.match_buffer(output, (c, h1 + h2, w), dtype=dtype)
|
|
|
|
for c_idx in T.thread_binding(c, thread="blockIdx.y"):
|
|
for h_idx, w_idx in T.grid(h1 + h2, w):
|
|
with T.sblock("dyn_concate_dim_1"):
|
|
T.reads(input_1_buf[c_idx, h_idx, w_idx])
|
|
T.writes(out_buf[c_idx, h_idx, w_idx])
|
|
if h_idx < h1:
|
|
out_buf[c_idx, h_idx, w_idx] = input_1_buf[c_idx, h_idx, w_idx]
|
|
else:
|
|
out_buf[c_idx, h_idx, w_idx] = input_2_buf[c_idx, h_idx - h1, w_idx]
|
|
|
|
return dyn_concate_dim_1_func
|
|
|
|
c1, h1, w1 = input_1.shape
|
|
c2, h2, w2 = input_2.shape
|
|
assert c1 == c2 and w1 == w2
|
|
|
|
out = op.tensor_ir_op(
|
|
create_dyn_concate_func(input_1.dtype),
|
|
"dyn_concate",
|
|
[input_1, input_2],
|
|
[Tensor.placeholder([c1, h1 + h2, w1], input_1.dtype)],
|
|
)
|
|
return out
|
|
|
|
def get_img_features(self, img_embeds: Tensor) -> Tensor:
|
|
img_processor_output = self.img_processor(img_embeds)
|
|
patch_feature = nn.op.split(img_processor_output, indices_or_sections=[1], axis=1)
|
|
return patch_feature[1]
|
|
|
|
def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
|
|
N, L, C = image_features.shape
|
|
num_images = 1
|
|
H = int(L**0.5)
|
|
image_features = nn.op.reshape(image_features, ([N, H, H, C])) # N, 24, 24, 1024
|
|
image_features = nn.op.reshape(
|
|
image_features, ([N, H // 2, 2, H // 2, 2, C])
|
|
) # N, 12, 2, 12, 2, 1024
|
|
|
|
new_s1 = tirx.Var("new_s1", "int64")
|
|
new_s2 = tirx.Var("new_s2", "int64")
|
|
|
|
image_features = op.wrap_nested(
|
|
relax.BlockBuilder()
|
|
.current()
|
|
.match_cast(
|
|
image_features._expr,
|
|
relax.TensorType(
|
|
[
|
|
image_features.shape[0],
|
|
new_s1,
|
|
image_features.shape[2],
|
|
new_s2,
|
|
image_features.shape[4],
|
|
image_features.shape[5],
|
|
],
|
|
image_features.dtype,
|
|
),
|
|
),
|
|
"image_features_1",
|
|
)
|
|
|
|
image_features = nn.op.permute_dims(
|
|
image_features, axes=([0, 1, 3, 2, 4, 5])
|
|
) # N, 12, 12, 2, 2, 1024
|
|
image_features = nn.op.reshape(image_features, ([N, -1, 4 * C])) # N, 144, 4096
|
|
image_features = nn.op.reshape(
|
|
image_features, ([num_images, h_crop, w_crop, H // 2, H // 2, -1])
|
|
)
|
|
|
|
new_s3 = tirx.Var("new_s3", "int64")
|
|
new_s4 = tirx.Var("new_s4", "int64")
|
|
|
|
image_features = op.wrap_nested(
|
|
relax.BlockBuilder()
|
|
.current()
|
|
.match_cast(
|
|
image_features._expr,
|
|
relax.TensorType(
|
|
[
|
|
image_features.shape[0],
|
|
new_s3,
|
|
image_features.shape[2],
|
|
new_s4,
|
|
image_features.shape[4],
|
|
image_features.shape[5],
|
|
],
|
|
image_features.dtype,
|
|
),
|
|
),
|
|
"image_features_2",
|
|
)
|
|
|
|
image_features = nn.op.permute_dims(image_features, axes=([0, 1, 3, 2, 4, 5]))
|
|
image_features_hd = nn.op.reshape(
|
|
image_features, ([num_images, h_crop * H // 2, w_crop * H // 2, 4 * C])
|
|
)
|
|
|
|
return image_features_hd
|
|
|
|
def add_image_newline(self, image_features_hd):
|
|
"""
|
|
image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
|
|
output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
|
|
"""
|
|
num_images, h, w, hid_dim = image_features_hd.shape
|
|
temp_sub_GN = self.dyn_repeat_4d_tensor(
|
|
self.sub_GN, T.int64(1), T.int64(h), T.int64(1), T.int64(1)
|
|
)
|
|
image_features_hd_newline = self.dyn_concate_dim_2(image_features_hd, temp_sub_GN)
|
|
image_features_hd_newline = nn.op.reshape(
|
|
image_features_hd_newline, ([num_images, -1, hid_dim])
|
|
)
|
|
return image_features_hd_newline
|
|
|
|
def forward(self, pixel_values: Tensor, h_crop, w_crop) -> Tensor:
|
|
img_features = self.get_img_features(pixel_values)
|
|
img_features = nn.op.split(img_features, indices_or_sections=[1], axis=0)
|
|
|
|
global_image_features = img_features[0]
|
|
global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
|
|
global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
|
|
|
|
sub_image_features = img_features[1]
|
|
sub_image_features_hd = self.reshape_hd_patches_2x2merge(sub_image_features, h_crop, w_crop)
|
|
sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
|
|
|
|
global_image_features_hd = nn.op.squeeze(global_image_features_hd_newline, 0)
|
|
|
|
combined_image = self.dyn_concate_dim_1(sub_image_features_hd_newline, self.glb_GN)
|
|
combined_image = self.dyn_concate_dim_1(combined_image, global_image_features_hd_newline)
|
|
combined_image = nn.op.squeeze(combined_image, 0)
|
|
|
|
new_s7 = tirx.Var("new_s7", "int64")
|
|
|
|
combined_image = op.wrap_nested(
|
|
relax.BlockBuilder()
|
|
.current()
|
|
.match_cast(
|
|
combined_image._expr,
|
|
relax.TensorType([new_s7, combined_image.shape[1]], combined_image.dtype),
|
|
),
|
|
"combined_image",
|
|
)
|
|
output_image = self.img_projection(combined_image)
|
|
return output_image
|