(tvm) 02/02: add handle add by byteoffset

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch refactor-s1
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit 4bd6ba1bcb4622eb79eff5ccca90ec077912d80b
Author: tqchen <ti...@gmail.com>
AuthorDate: Mon Apr 14 10:50:26 2025 -0400

    add handle add by byteoffset
---
 include/tvm/tir/builtin.h              | 11 +++++++++++
 python/tvm/script/ir_builder/tir/ir.py |  2 ++
 python/tvm/tir/__init__.py             |  2 +-
 python/tvm/tir/op.py                   | 19 +++++++++++++++++++
 src/tir/op/builtin.cc                  |  4 ++++
 5 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h
index 2579250de2..822763d0b2 100644
--- a/include/tvm/tir/builtin.h
+++ b/include/tvm/tir/builtin.h
@@ -263,6 +263,15 @@ TVM_DLL const Op& tvm_context_id();
  */
 TVM_DLL const Op& tvm_tuple();
 
+/*!
+ * \brief See pesudo code
+ *
+ *  void* handle_add_byte_offset(void* handle, int offset) {
+ *     return reinterpret_cast<v*>(reinterpret_cast<char*>(handle) + offset);
+ *  }
+ */
+TVM_DLL const Op& handle_add_byte_offset();
+
 /*!
  * \brief See pesudo code
  *
@@ -990,6 +999,8 @@ enum TVMStructFieldKind : int {
   kArrKindBound_,
   // TVMValue field
   kTVMValueContent,
+  kTVMFFIAnyTypeIndex,
+  kTVMFFIAnyUnionValue,
   kTVMValueKindBound_
 };
 }  // namespace builtin
diff --git a/python/tvm/script/ir_builder/tir/ir.py b/python/tvm/script/ir_builder/tir/ir.py
index 3e835e8d9d..3814f2df88 100644
--- a/python/tvm/script/ir_builder/tir/ir.py
+++ b/python/tvm/script/ir_builder/tir/ir.py
@@ -1881,6 +1881,7 @@ call_cpacked = _op_wrapper(_tir_op.call_cpacked)
 call_packed_lowered = _op_wrapper(_tir_op.call_packed_lowered)
 call_cpacked_lowered = _op_wrapper(_tir_op.call_cpacked_lowered)
 tvm_tuple = _op_wrapper(_tir_op.tvm_tuple)
+handle_add_byte_offset = _op_wrapper(_tir_op.handle_add_byte_offset)
 tvm_struct_set = _op_wrapper(_tir_op.tvm_struct_set)
 tvm_struct_get = _tir_op.tvm_struct_get
 tvm_thread_invariant = _op_wrapper(_tir_op.tvm_thread_invariant)
@@ -2178,6 +2179,7 @@ __all__ = [
     "call_llvm_pure_intrin",
     "call_pure_extern",
     "tvm_tuple",
+    "handle_add_byte_offset",
     "tvm_struct_set",
     "tvm_struct_get",
     "tvm_thread_invariant",
diff --git a/python/tvm/tir/__init__.py b/python/tvm/tir/__init__.py
index 4f56ec3c15..63d3cb8f31 100644
--- a/python/tvm/tir/__init__.py
+++ b/python/tvm/tir/__init__.py
@@ -50,7 +50,7 @@ from .op import call_packed, call_cpacked, call_intrin, call_pure_extern, call_e
 from .op import call_llvm_intrin, call_llvm_pure_intrin, ret, all, any, min_value, max_value, trace
 from .op import tvm_check_return
 from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
-from .op import tvm_tuple, tvm_struct_get, tvm_struct_set
+from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get, tvm_struct_set
 from .op import address_of, lookup_param, assume, undef
 from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr, tvm_throw_last_error
 from .op import (
diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
index 53c92fff86..94865bcd29 100644
--- a/python/tvm/tir/op.py
+++ b/python/tvm/tir/op.py
@@ -495,6 +495,25 @@ def tvm_tuple(*value):
     return call_intrin("handle", "tir.tvm_tuple", *value)
 
 
+def handle_add_offset(handle, offset):
+    """Add offset to handle
+
+    Parameters
+    ----------
+    handle : Expr
+        The handle.
+
+    offset : int
+        The offset.
+
+    Returns
+    -------
+    call : PrimExpr
+        The call expression.
+    """
+    return call_intrin("handle", "tir.handle_add_offset", handle, offset)
+
+
 def tvm_struct_get(arr, index, field, dtype):
     """Get struct field value in array
 
diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc
index 39a09d9922..f0894acecf 100644
--- a/src/tir/op/builtin.cc
+++ b/src/tir/op/builtin.cc
@@ -163,6 +163,10 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_context_id)
 TIR_DEFINE_BUILTIN_FUNC(tvm_tuple).set_attr<TCallEffectKind>("TCallEffectKind",
                                                              Integer(CallEffectKind::kEmbedInfo));
 
+TIR_DEFINE_BUILTIN_FUNC(handle_add_byte_offset)
+    .set_num_inputs(2)
+    .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
+
 TIR_DEFINE_BUILTIN_FUNC(tvm_struct_get)
     .set_num_inputs(3)
     .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState))