[TIR] Expose Stack-related TVM builtins in Python

python/tvm/tir/__init__.py

@@ -47,6 +47,7 @@
 
 from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_extern
 from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any, min_value, max_value, trace
+from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
 from .op import exp, exp2, exp10, log, log2, log10, log1p, ldexp, clz
 from .op import sin, sinh, asin, asinh
 from .op import cos, cosh, acos, acosh

python/tvm/tir/op.py

@@ -250,6 +250,74 @@ def call_llvm_pure_intrin(dtype, name, *args, span=None):
     )
 
 
+def tvm_stack_alloca(dtype_str, num):
+    """Return new on stack dtype[num]
+
+    Parameters
+    ----------
+    dtype_str : str
+        The data type of array.
+
+    num : int
+        The size of array.
+
+    Returns
+    -------
+    call : PrimExpr
+        The call expression.
+    """
+    return call_intrin("handle", "tir.tvm_stack_alloca", dtype_str, num)
+
+
+def tvm_stack_make_shape(*args):
+    """Allocate a shape tuple on stack, return the handle
+
+    Parameters
+    ----------
+    args : int
+        The tuple shape.
+
+    Returns
+    -------
+    call : PrimExpr
+        The call expression.
+    """
+    return call_intrin("handle", "tir.tvm_stack_make_shape", *args)
+
+
+def tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset):
+    """Allocate a NDArray(DLTensor) on stack, return the handle
+
+    Parameters
+    ----------
+    data : Expr
+        The data of array.
+
+    shape : Expr
+        The shape of array.
+
+    strides : Expr
+        The strides of array.
+
+    ndim : Expr
+        The dimensions of array.
+
+    arr_dtype : Expr
+        The data type of array.
+
+    elem_offse : Expr
+        The element offset of array.
+
+    Returns
+    -------
+    call : PrimExpr
+        The call expression.
+    """
+    return call_intrin(
+        "handle", "tir.tvm_stack_make_array", data, shape, strides, ndim, arr_dtype, elem_offset
+    )
+
+
 def ret(val):
     """Create a tir return expression