Skip to content

[Web] Decode packed BF16 tensor records in place #20167

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akaashrp wants to merge 1 commit into
base: main
Choose a base branch
from
+566 −29
Open

[Web] Decode packed BF16 tensor records in place #20167

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 73 additions & 16 deletions web/emcc/wasm_runtime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>

#include <limits>

#include "src/runtime/cpu_device_api.cc"
#include "src/runtime/device_api.cc"
#include "src/runtime/extra/contrib/sort/sort.cc"
Expand Down Expand Up @@ -126,38 +128,93 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}

size_t GetCheckedTensorElementCount(const Tensor& tensor) {
size_t size = 1;
for (int i = 0; i < tensor->ndim; ++i) {
TVM_FFI_ICHECK_GE(tensor->shape[i], 0);
const uint64_t dim_u64 = static_cast<uint64_t>(tensor->shape[i]);
TVM_FFI_ICHECK_LE(dim_u64, std::numeric_limits<size_t>::max());
const size_t dim = static_cast<size_t>(dim_u64);
if (size != 0) {
TVM_FFI_ICHECK_LE(dim, std::numeric_limits<size_t>::max() / size);
}
size *= dim;
}
return size;
}

void CheckF32CPUTensor(const Tensor& tensor) {
TVM_FFI_ICHECK_EQ(tensor->device.device_type, kDLCPU);
TVM_FFI_ICHECK(tensor.IsContiguous());
TVM_FFI_ICHECK_EQ(tensor->dtype.code, kDLFloat);
TVM_FFI_ICHECK_EQ(tensor->dtype.bits, 32);
TVM_FFI_ICHECK_EQ(tensor->dtype.lanes, 1);
}

void ArrayDecodeStorage(Tensor cpu_arr, TVMFFIByteArray* bytes, const std::string& format,
const std::string& dtype) {
TVM_FFI_ICHECK_NE(bytes, nullptr);
const char* byte_data = bytes->data;
const size_t byte_size = bytes->size;
if (format == "f32-to-bf16" && dtype == "float32") {
const uint16_t* bf16 = reinterpret_cast<const uint16_t*>(byte_data);
uint32_t* data = static_cast<uint32_t*>(cpu_arr->data);
TVM_FFI_ICHECK(cpu_arr.IsContiguous());
size_t size = 1;
for (int i = 0; i < cpu_arr->ndim; ++i) {
size *= cpu_arr->shape[i];
CheckF32CPUTensor(cpu_arr);
const size_t size = GetCheckedTensorElementCount(cpu_arr);
TVM_FFI_ICHECK_LE(size, std::numeric_limits<size_t>::max() / 4);
TVM_FFI_ICHECK_EQ(byte_size, size * 2);
if (size == 0) {
return;
}
TVM_FFI_ICHECK_EQ(size, byte_size / 2);
TVM_FFI_ICHECK_NE(cpu_arr->data, nullptr);
TVM_FFI_ICHECK_NE(byte_data, nullptr);
const uint8_t* bf16 = reinterpret_cast<const uint8_t*>(byte_data);
uint8_t* data = static_cast<uint8_t*>(cpu_arr->data) + cpu_arr->byte_offset;
for (size_t i = 0; i < size; ++i) {
data[i] = static_cast<uint32_t>(bf16[i]) << 16;
data[4 * i] = 0;
data[4 * i + 1] = 0;
data[4 * i + 2] = bf16[2 * i];
data[4 * i + 3] = bf16[2 * i + 1];
}
} else {
cpu_arr.CopyFromBytes(byte_data, byte_size);
}
}

void ArrayDecodeBF16ToF32Inplace(Tensor cpu_arr, int64_t encoded_nbytes) {
CheckF32CPUTensor(cpu_arr);
TVM_FFI_ICHECK_GE(encoded_nbytes, 0);

const size_t size = GetCheckedTensorElementCount(cpu_arr);
TVM_FFI_ICHECK_LE(size, std::numeric_limits<size_t>::max() / 4);
TVM_FFI_ICHECK_EQ(static_cast<uint64_t>(encoded_nbytes), static_cast<uint64_t>(size) * 2);
if (size == 0) {
return;
}
TVM_FFI_ICHECK_NE(cpu_arr->data, nullptr);

uint8_t* data = static_cast<uint8_t*>(cpu_arr->data) + cpu_arr->byte_offset;
for (size_t i = size; i != 0; --i) {
const size_t j = i - 1;
const uint8_t low = data[2 * j];
const uint8_t high = data[2 * j + 1];
data[4 * j] = 0;
data[4 * j + 1] = 0;
data[4 * j + 2] = low;
data[4 * j + 3] = high;
}
}

TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def_packed(
"tvmjs.array.decode_storage", [](ffi::PackedArgs args, ffi::Any* ret) {
Tensor cpu_arr = args[0].cast<Tensor>();
TVMFFIByteArray* bytes = args[1].cast<TVMFFIByteArray*>();
std::string format = args[2].cast<ffi::String>().operator std::string();
std::string dtype = args[3].cast<ffi::String>().operator std::string();
ArrayDecodeStorage(cpu_arr, bytes, format, dtype);
});
refl::GlobalDef()
.def_packed("tvmjs.array.decode_storage",
[](ffi::PackedArgs args, ffi::Any* ret) {
Tensor cpu_arr = args[0].cast<Tensor>();
TVMFFIByteArray* bytes = args[1].cast<TVMFFIByteArray*>();
std::string format = args[2].cast<ffi::String>().operator std::string();
std::string dtype = args[3].cast<ffi::String>().operator std::string();
ArrayDecodeStorage(cpu_arr, bytes, format, dtype);
})
.def("tvmjs.array.decode_bf16_to_f32_inplace", ArrayDecodeBF16ToF32Inplace);
}

// Concatenate n TVMArrays
Expand Down
89 changes: 80 additions & 9 deletions web/src/runtime.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@ class RuntimeContext implements Disposable {
tensorCacheRemove: PackedFunc;
tensorCacheClear: PackedFunc;
arrayDecodeStorage: PackedFunc;
arrayDecodeBF16ToF32Inplace: PackedFunc | undefined;
paramModuleFromCache: PackedFunc;
paramModuleFromCacheByName: PackedFunc;
makeShapeTuple: PackedFunc;
Expand Down Expand Up @@ -214,6 +215,7 @@ class RuntimeContext implements Disposable {
this.tensorCacheUpdate = getGlobalFunc("vm.builtin.tensor_cache.update");
this.tensorCacheClear = getGlobalFunc("vm.builtin.tensor_cache.clear");
this.arrayDecodeStorage = getGlobalFunc("tvmjs.array.decode_storage");
this.arrayDecodeBF16ToF32Inplace = undefined;
this.paramModuleFromCache = getGlobalFunc("vm.builtin.param_module_from_cache");
this.paramModuleFromCacheByName = getGlobalFunc("vm.builtin.param_module_from_cache_by_name");
this.makeShapeTuple = getGlobalFunc("ffi.Shape");
Expand All @@ -237,6 +239,7 @@ class RuntimeContext implements Disposable {
this.tensorCacheRemove.dispose();
this.tensorCacheUpdate.dispose();
this.arrayDecodeStorage.dispose();
this.arrayDecodeBF16ToF32Inplace?.dispose();
this.paramModuleFromCache.dispose();
this.paramModuleFromCacheByName.dispose();
this.makeShapeTuple.dispose();
Expand Down Expand Up @@ -606,6 +609,22 @@ export class Tensor extends TVMObject {
return this.dataPtr;
}

/**
* Return the effective address of a CPU Tensor's storage.
* @returns The address in Wasm linear memory.
* @internal
*/
getCPUDataAddress(): Pointer {
if (this.device.deviceType !== DeviceStrToEnum.cpu) {
throw new Error("Can only obtain a linear-memory address for a CPU Tensor");
}
const address = this.getDataPtr() + this.byteOffset;
if (!Number.isSafeInteger(address) || address < 0) {
throw new Error("Invalid CPU Tensor storage address");
}
return address;
}

/**
* Copy data from another Tensor or javascript array.
* The number of elements must match.
Expand Down Expand Up @@ -953,6 +972,10 @@ export class Instance implements Disposable {
return this.getGlobalFuncInternal(name, autoAttachToScope);
}
);
this.ctx.arrayDecodeBF16ToF32Inplace = this.getGlobalFuncInternalOptional(
"tvmjs.array.decode_bf16_to_f32_inplace",
/*autoAttachToScope=*/ false,
);
this.registerEnvGlobalPackedFuncs();
this.registerObjectFactoryFuncs();
this.rng = new LinearCongruentialGenerator();
Expand Down Expand Up @@ -1156,6 +1179,17 @@ export class Instance implements Disposable {
}

private getGlobalFuncInternal(name: string, autoAttachToScope = true): PackedFunc {
const ret = this.getGlobalFuncInternalOptional(name, autoAttachToScope);
if (ret === undefined) {
throw Error("Cannot find global function " + name);
}
return ret;
}

private getGlobalFuncInternalOptional(
name: string,
autoAttachToScope = true,
): PackedFunc | undefined {
const stack = this.lib.getOrAllocCallStack();
const nameOffset = stack.allocByteArrayForString(name);
const outOffset = stack.allocPtrArray(1);
Expand All @@ -1172,7 +1206,7 @@ export class Instance implements Disposable {
const handle = this.memory.loadPointer(outPtr);
this.lib.recycleCallStack(stack);
if (handle === 0) {
throw Error("Cannot find global function " + name);
return undefined;
}
const ret = this.makePackedFunc(handle);
if (autoAttachToScope) this.ctx.attachToCurrentScope(ret);
Expand Down Expand Up @@ -1357,12 +1391,22 @@ export class Instance implements Disposable {
): void {
const recBytes = getTensorCacheRecordBytes(shardBytes, rec);
if (cpuArray !== undefined) {
this.ctx.arrayDecodeStorage(
cpuArray,
recBytes,
rec.format,
rec.dtype,
);
const isPackedBF16 =
rec.format === "f32-to-bf16" && rec.dtype === "float32";
if (isPackedBF16 && this.ctx.arrayDecodeBF16ToF32Inplace !== undefined) {
this.memory.storeRawBytes(cpuArray.getCPUDataAddress(), recBytes);
this.ctx.arrayDecodeBF16ToF32Inplace(
cpuArray,
new Scalar(recBytes.byteLength, "int64"),
);
} else {
this.ctx.arrayDecodeStorage(
cpuArray,
recBytes,
rec.format,
rec.dtype,
);
}
}
if (gpuArray !== undefined) {
if (cpuArray === undefined) {
Expand All @@ -1373,6 +1417,24 @@ export class Instance implements Disposable {
}
}

/** Return the exact byte size of a packed BF16 tensor. */
private getPackedBF16Bytes(shape: Array<number>): number {
let numElements = 1;
for (const dim of shape) {
if (!Number.isSafeInteger(dim) || dim < 0) {
throw new Error(`Invalid tensor dimension: ${dim}`);
}
numElements *= dim;
if (!Number.isSafeInteger(numElements)) {
throw new Error("Tensor element count exceeds JavaScript's safe integer range");
}
}
const nbytes = numElements * 2;
if (!Number.isSafeInteger(nbytes)) {
throw new Error("Packed BF16 size exceeds JavaScript's safe integer range");
}
return nbytes;
}

/**
* Fetch list of Tensor into the TensorCache.
Expand Down Expand Up @@ -1487,11 +1549,20 @@ export class Instance implements Disposable {
let gpu_arr: Tensor | undefined;
try {
const rec = shardRecords[j];
const requiresDecode =
const isPackedBF16 =
rec.format === "f32-to-bf16" && rec.dtype === "float32";
if (isPackedBF16) {
const expectedBytes = this.getPackedBF16Bytes(rec.shape);
if (rec.nbytes !== expectedBytes) {
throw new Error(
`Packed BF16 record has ${rec.nbytes} bytes, ` +
`but shape requires ${expectedBytes}`,
);
}
}
const directToWebGPU =
device.deviceType === DeviceStrToEnum.webgpu &&
!requiresDecode &&
!isPackedBF16 &&
rec.nbytes % 4 === 0;

if (!directToWebGPU) {
Expand Down
Loading
Loading