chore: import upstream snapshot with attribution

python/tvm/relax/op/vm/__init__.py

@@ -0,0 +1,19 @@
+# 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 vm primitives."""
+
+from .vm import alloc_storage, alloc_tensor, call_tir_dyn, kill_object

python/tvm/relax/op/vm/_ffi_api.py

@@ -0,0 +1,20 @@
+# 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
+"""FFI APIs for tvm.relax.op.vm"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.op.vm", __name__)

python/tvm/relax/op/vm/vm.py

@@ -0,0 +1,144 @@
+# 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 vm primitives."""
+
+from tvm.ir import Call
+
+from ...expr import DataTypeImm, Expr, StringImm, Tuple, prim_value
+from ...utils import convert_to_expr
+from . import _ffi_api
+
+
+def alloc_storage(
+    shape: Expr,
+    runtime_device_index: int | Expr,
+    dtype: str | Expr,
+    storage_scope: str | StringImm = "global",
+) -> Call:
+    """Construct a Call to allocate a storage with specific size,
+    runtime_device_index, and dtype.
+
+    Parameters
+    ----------
+    shape : Expr
+        The shape of the storage to be allocated.
+
+    runtime_device_index : 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.
+
+    dtype : Union[str, Expr]
+        The datatype of the storage to be allocated.
+
+    storage_scope : Union[str, StringImm]
+        The storage scope of the storage to allocate. Default is global.
+
+    Returns
+    -------
+    result : Call
+        A relax Call, which gets the allocated storage.
+    """
+    shape = convert_to_expr(shape)
+    if isinstance(dtype, str):
+        dtype = DataTypeImm(dtype)
+    if isinstance(storage_scope, str):
+        storage_scope = StringImm(storage_scope)
+    if isinstance(runtime_device_index, int):
+        runtime_device_index = prim_value(runtime_device_index)
+    return _ffi_api.alloc_storage(shape, runtime_device_index, dtype, storage_scope)  # 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_object(obj: Expr) -> Call:
+    """Construct a Call to set the register corresponding to the input object to
+    null at runtime, in order to kill the input object.
+
+    Parameters
+    ----------
+    obj : Expr
+        The object to be killed.
+
+    Returns
+    -------
+    result : Call
+        CallNode that kills the input object.
+    """
+    return _ffi_api.kill_object(obj)  # type: ignore
+
+
+def call_tir_dyn(func: Expr, args: Tuple) -> Call:
+    """Construct a Call to call_tir_dyn (invoke the given TIR PrimFunc)
+    consisting of the input tensors and the shape of the result.
+
+    Parameters
+    ----------
+    func : Expr
+        An expression evaluating to a TIR PrimFunc.
+
+    args : Tuple
+        The input args, includes a list of tensors, and a ShapeExpr.
+
+    Returns
+    -------
+    result : Call
+        A relax Call to call_tir_dyn.
+    """
+    func = convert_to_expr(func)
+    if isinstance(args, list | tuple):
+        args = Tuple(args)
+
+    return _ffi_api.call_tir_dyn(func, args)  # type: ignore