wehub-resource-sync
136 lines
4.0 KiB
Python
136 lines
4.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
|
|
"""Relax memory primitives."""
|
|
|
|
from tvm.ir import Call
|
|
|
|
from ...expr import DataTypeImm, Expr, StringImm, prim_value
|
|
from ...utils import convert_to_expr
|
|
from . import _ffi_api
|
|
|
|
|
|
def alloc_storage(
|
|
size: Expr,
|
|
virtual_device_index: int | Expr,
|
|
storage_scope: str | Expr,
|
|
dtype: str | Expr,
|
|
) -> Call:
|
|
"""Construct a Call to allocate a storage with specific size, virtual_device_index,
|
|
storage_scope and dtype.
|
|
|
|
Parameters
|
|
----------
|
|
size : Expr
|
|
The size of the storage to be allocated.
|
|
|
|
virtual_device_index : Union[int, Expr]
|
|
The virtual device index indicating on which device the storage is to be allocated.
|
|
Index -1 is reserved for the host device.
|
|
|
|
storage_scope : Union[str, Expr]
|
|
The storage scope to allocate the storage to.
|
|
|
|
dtype : Union[str, Expr]
|
|
The datatype of the storage to be allocated.
|
|
|
|
Returns
|
|
-------
|
|
result : Call
|
|
A relax Call, which gets the allocated storage.
|
|
"""
|
|
size = convert_to_expr(size)
|
|
if isinstance(dtype, str):
|
|
dtype = DataTypeImm(dtype)
|
|
if isinstance(storage_scope, str):
|
|
storage_scope = StringImm(storage_scope)
|
|
if isinstance(virtual_device_index, int):
|
|
virtual_device_index = prim_value(virtual_device_index)
|
|
return _ffi_api.alloc_storage(size, virtual_device_index, storage_scope, dtype) # type: ignore
|
|
|
|
|
|
def alloc_tensor(
|
|
storage: Expr,
|
|
offset: int | Expr,
|
|
shape: Expr,
|
|
dtype: str | Expr,
|
|
runtime_device_ind: int | Expr = prim_value(0),
|
|
) -> Call:
|
|
"""Construct a Call to allocate a tensor on a certain storage starting from the given offset.
|
|
|
|
Parameters
|
|
----------
|
|
storage : Expr
|
|
The storage to allocate the tensor to.
|
|
|
|
offset : Union[int, Expr]
|
|
The storage offset to allocate the tensor.
|
|
|
|
shape : Expr
|
|
The shape of the tensor to be allocated.
|
|
|
|
dtype : Union[str, Expr]
|
|
The datatype of the tensor to be allocated.
|
|
|
|
runtime_device_ind: Union[int, Expr]
|
|
The device index indicating on which device the tensor is to be
|
|
allocated at runtime. Index -1 is reserved for the host device.
|
|
|
|
Returns
|
|
-------
|
|
result : Call
|
|
A relax Call, which gets the allocated tensor.
|
|
"""
|
|
if isinstance(offset, int):
|
|
offset = prim_value(offset)
|
|
shape = convert_to_expr(shape)
|
|
if isinstance(dtype, str):
|
|
dtype = DataTypeImm(dtype)
|
|
if isinstance(runtime_device_ind, int):
|
|
runtime_device_ind = prim_value(runtime_device_ind)
|
|
return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind) # type: ignore
|
|
|
|
|
|
def kill_storage(storage: Expr) -> Call:
|
|
"""Construct a Call to kill a storage.
|
|
|
|
Parameters
|
|
----------
|
|
storage : Expr
|
|
The storage to be killed.
|
|
|
|
Returns
|
|
-------
|
|
result : Call
|
|
A relax Call to kill a storage.
|
|
"""
|
|
return _ffi_api.kill_storage(storage) # type: ignore
|
|
|
|
|
|
def kill_tensor(tensor: Expr) -> Call:
|
|
"""Construct a Call to kill a tensor.
|
|
|
|
Parameters
|
|
----------
|
|
tensor : Expr
|
|
The tensor to be killed.
|
|
|
|
Returns
|
|
-------
|
|
result : Call
|
|
A relax Call to kill a tensor.
|
|
"""
|
|
return _ffi_api.kill_tensor(tensor) # type: ignore
|