This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvm-further-cleanup-python-tests-followup in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 72368f247354b338121799c42cbf66914f432aeb Author: Tianqi Chen <[email protected]> AuthorDate: Sun Jul 5 18:46:38 2026 +0000 [FIX][TE] Zero offsets for internal extern buffers --- src/te/operation/create_primfunc.cc | 11 +++++++-- tests/python/te/test_te_create_primfunc.py | 37 ++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc index 3d1b0fa0d3..d64b768e0e 100644 --- a/src/te/operation/create_primfunc.cc +++ b/src/te/operation/create_primfunc.cc @@ -641,10 +641,17 @@ Stmt GenerateStmtFromExternOp(const te::ExternOp& extern_op, CreateFuncInfo* inf for (int i = 0; i < extern_op->num_outputs(); ++i) { const Buffer& placeholder = extern_op->output_placeholders[i]; const te::Tensor& output_tensor = extern_op.output(i); - info->tensor2buffers[output_tensor] = placeholder; + Buffer output_buffer = placeholder; if (!info->IsArg(output_tensor)) { - info->root_alloc.push_back(placeholder); + PrimExpr zero_offset = IntImm(placeholder->elem_offset.ty(), 0); + if (const auto* offset_var = placeholder->elem_offset.as<VarNode>()) { + var_map[offset_var] = zero_offset; + } + output_buffer.CopyOnWrite()->elem_offset = zero_offset; + input_buffer_map[placeholder.get()] = output_buffer; + info->root_alloc.push_back(output_buffer); } + info->tensor2buffers[output_tensor] = output_buffer; } // The access region does not need to be collected here, as it will diff --git a/tests/python/te/test_te_create_primfunc.py b/tests/python/te/test_te_create_primfunc.py index 90d7b69a55..48f9c1bf7b 100644 --- a/tests/python/te/test_te_create_primfunc.py +++ b/tests/python/te/test_te_create_primfunc.py @@ -293,6 +293,43 @@ def test_extern(): _check_workload(te_extern, tir_extern) +def te_extern_epilogue(): + A = te.placeholder((4, 3), name="A") + B = te.placeholder((3, 2), name="B") + C = te.extern( + (4, 2), + [A, B], + lambda ins, outs: tvm.tirx.call_packed("testing.echo", ins[0], ins[1], outs[0]), + name="C", + ) + D = te.compute(C.shape, lambda i, j: C[i, j] + 1.0, name="D") + return [A, B, D] + + [email protected]_func(s_tir=True) +def tir_extern_epilogue(var_A: T.handle, var_B: T.handle, D: T.Buffer((4, 2), "float32")): + T.func_attr({"global_symbol": "main", "tirx.noalias": True}) + A = T.match_buffer(var_A, (4, 3), offset_factor=1) + B = T.match_buffer(var_B, (3, 2), offset_factor=1) + C = T.sblock_alloc_buffer((4, 2), elem_offset=0, offset_factor=1) + with T.sblock("C"): + T.reads() + T.writes() + T.call_packed("testing.echo", A, B, C) + for i, j in T.grid(4, 2): + with T.sblock("D"): + vi, vj = T.axis.remap("SS", [i, j]) + T.reads(C[vi, vj]) + T.writes(D[vi, vj]) + D[vi, vj] = C[vi, vj] + T.float32(1) + + +def test_extern_epilogue(): + _check_workload(te_extern_epilogue, tir_extern_epilogue) + func = te.create_prim_func(te_extern_epilogue()).with_attr("global_symbol", "extern_epilogue") + tvm.compile(func, target="llvm") + + def te_reordered_matmul(): k = te.reduce_axis((0, 128), "k") A = te.placeholder((128, 128), name="A")
