wehub-resource-sync
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
"""The pass that attaches embedding allocation function to the IRModule."""
|
|
|
|
from typing import Any, Dict # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import IRModule, relax
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="AttachAllocEmbeddingTensorFunc")
|
|
class AttachAllocEmbeddingTensorFunc:
|
|
"""Attach embedding tensor allocation Relax function to IRModule."""
|
|
|
|
def __init__(self, metadata: Dict[str, Any]): # noqa: UP006
|
|
self.metadata = metadata
|
|
|
|
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
|
"""Entrypoint"""
|
|
embed_func = None
|
|
for gv, func in mod.functions_items():
|
|
if gv.name_hint == "embed":
|
|
embed_func = func
|
|
|
|
if embed_func is None:
|
|
return mod
|
|
|
|
hidden_size = embed_func.ret_ty.shape[-1]
|
|
dtype = relax.DataTypeImm(embed_func.ret_ty.dtype.dtype)
|
|
bb = relax.BlockBuilder(mod)
|
|
with bb.function("alloc_embedding_tensor", []):
|
|
bb.emit_func_output(
|
|
bb.emit(
|
|
relax.op.builtin.alloc_tensor(
|
|
relax.ShapeExpr([self.metadata["prefill_chunk_size"], hidden_size]),
|
|
dtype,
|
|
runtime_device_index=0,
|
|
)
|
|
)
|
|
)
|
|
return bb.finalize()
|