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 9ee70e95dc88c3a29d4cc61e37667610a6f47b87 Author: tqchen <[email protected]> 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 | 91 ++++++++++++++++++++++++++++------ src/tir/op/builtin.cc | 4 ++ 5 files changed, 95 insertions(+), 15 deletions(-) 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..2e0d2a64c9 100644 --- a/python/tvm/tir/op.py +++ b/python/tvm/tir/op.py @@ -32,7 +32,11 @@ from .expr import Call, CommReducer, IntImm, PrimExprWithOp, Var def _pack_buffer(buf, span=None): """Build intrinsics that packs the buffer.""" shape = Call("handle", "tir.tvm_stack_make_shape", buf.shape, span) - strides = Call("handle", "tir.tvm_stack_make_shape", buf.strides, span) if buf.strides else 0 + strides = ( + Call("handle", "tir.tvm_stack_make_shape", buf.strides, span) + if buf.strides + else 0 + ) pack_args = [ buf.data, shape, @@ -334,7 +338,9 @@ def tvm_check_return(expected, return_unexpected, nested_call): call : PrimExpr The call expression. """ - return call_intrin("int32", "tir.tvm_check_return", expected, return_unexpected, nested_call) + return call_intrin( + "int32", "tir.tvm_check_return", expected, return_unexpected, nested_call + ) def tvm_stack_alloca(dtype_str, num): @@ -401,7 +407,14 @@ def tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset): The call expression. """ return call_intrin( - "handle", "tir.tvm_stack_make_array", data, shape, strides, ndim, arr_dtype, elem_offset + "handle", + "tir.tvm_stack_make_array", + data, + shape, + strides, + ndim, + arr_dtype, + elem_offset, ) @@ -495,6 +508,25 @@ def tvm_tuple(*value): return call_intrin("handle", "tir.tvm_tuple", *value) +def handle_add_byte_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_byte_offset", handle, offset) + + def tvm_struct_get(arr, index, field, dtype): """Get struct field value in array @@ -653,7 +685,9 @@ def tvm_warp_shuffle(mask, value, warp_id, width, warp_size): call : PrimExpr The call expression. """ - return call_intrin(value.dtype, "tir.tvm_warp_shuffle", mask, value, warp_id, width, warp_size) + return call_intrin( + value.dtype, "tir.tvm_warp_shuffle", mask, value, warp_id, width, warp_size + ) def tvm_warp_shuffle_up(mask, value, offset, width, warp_size): @@ -762,7 +796,9 @@ def tvm_access_ptr(ptype, data, offset, extent, rw_mask): call : PrimExpr The call expression. """ - return call_intrin("handle", "tir.tvm_access_ptr", ptype, data, offset, extent, rw_mask) + return call_intrin( + "handle", "tir.tvm_access_ptr", ptype, data, offset, extent, rw_mask + ) def tvm_throw_last_error(): @@ -1298,7 +1334,9 @@ def mma_fill(dtype, local_size, local_ptr, offset): ) -def ptx_ldmatrix(dtype, trans, num, type, local_ptr, local_offset, smem_ptr, smem_offset): +def ptx_ldmatrix( + dtype, trans, num, type, local_ptr, local_offset, smem_ptr, smem_offset +): """TVM intrinsic for ptx load matrix from shared memory https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix @@ -1376,7 +1414,13 @@ def ptx_cp_async(dtype, shared_ptr, shared_offset, global_ptr, global_offset, by The call expression. """ return call_intrin( - dtype, "tir.ptx_cp_async", shared_ptr, shared_offset, global_ptr, global_offset, bytes + dtype, + "tir.ptx_cp_async", + shared_ptr, + shared_offset, + global_ptr, + global_offset, + bytes, ) @@ -1489,7 +1533,9 @@ def ptx_init_barrier_thread_count(barrier_id, thread_count): call : PrimExpr The call expression. """ - return call_intrin("", "tir.ptx_init_barrier_thread_count", barrier_id, thread_count) + return call_intrin( + "", "tir.ptx_init_barrier_thread_count", barrier_id, thread_count + ) def ptx_arrive_barrier(barrier_id): @@ -1595,7 +1641,9 @@ def make_filled_simdgroup_matrix( call : PrimExpr The call expression. """ - return call_intrin("handle", "tir.make_filled_simdgroup_matrix", d, index, value, col, row) + return call_intrin( + "handle", "tir.make_filled_simdgroup_matrix", d, index, value, col, row + ) def simdgroup_load( @@ -1691,7 +1739,15 @@ def simdgroup_store( The call expression. """ return call_intrin( - "handle", "tir.simdgroup_store", d, index, ptr, stride, col, row, transpose_matrix + "handle", + "tir.simdgroup_store", + d, + index, + ptr, + stride, + col, + row, + transpose_matrix, ) @@ -3348,10 +3404,13 @@ def comm_reducer(fcombine, fidentity, name="reduce"): if where is None: where = tir.convert(True) if init is None: - outputs = tuple(tvm.tir.Reduce(combiner, expr, axis, where, i, []) for i in range(size)) + outputs = tuple( + tvm.tir.Reduce(combiner, expr, axis, where, i, []) for i in range(size) + ) else: outputs = tuple( - tvm.tir.Reduce(combiner, expr, axis, where, i, init) for i in range(size) + tvm.tir.Reduce(combiner, expr, axis, where, i, init) + for i in range(size) ) return outputs[0] if size == 1 else outputs @@ -3407,7 +3466,9 @@ def comm_reducer(fcombine, fidentity, name="reduce"): return reducer -def TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint): +def TVMBackendAllocWorkspace( + device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint +): """Backend function to allocate temporal workspace Parameters @@ -3462,7 +3523,9 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr): call : PrimExpr The call expression. """ - return call_intrin("int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr) + return call_intrin( + "int32", "tir.TVMBackendFreeWorkspace", device_type, device_id, ptr + ) def anylist_getitem(list_handle, index): 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))
