This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 53cee4bca3 [TVMScript] Round-trip DeclBuffer with undefined data
pointer (#14900)
53cee4bca3 is described below
commit 53cee4bca3743a2e26bcb1bfe7aa2aeb3956040d
Author: Eric Lunderberg <[email protected]>
AuthorDate: Sun May 21 22:15:51 2023 -0500
[TVMScript] Round-trip DeclBuffer with undefined data pointer (#14900)
Previously, a `DeclBuffer` object in which `decl_buffer->buffer->data`
had not been defined would be printed without the `data` argument.
This was the same representation as the `Allocate`/`DeclBuffer`
pattern (e.g. `buf = T.decl_buffer(shape,dtype)`), and so the
`Allocate` node would be inserted when parsing.
This commit updates the printing of buffers to be aware of which
variables are implicitly declared in the process of declaring a
buffer. (e.g. A buffer that occurs in `BlockNode::alloc_buffers`
defines both the buffer and the data pointer, while a `DeclBuffer`
only defines the buffer.) This is used to produce a different
representation for the undefined data pointer, removing the ambiguity
with the `Allocate`/`DeclBuffer` sugar.
Because this change only affects malformed TIR in which a `DeclBuffer`
node references an undefined data pointer, it is intended primarily
for debugging purposes.
---
src/script/printer/tir/block.cc | 3 +-
src/script/printer/tir/buffer.cc | 29 +++++++-----
src/script/printer/tir/function.cc | 6 ++-
src/script/printer/tir/stmt.cc | 24 ++++++----
src/script/printer/tir/utils.h | 28 +++++++++++-
.../python/unittest/test_tvmscript_printer_tir.py | 3 +-
tests/python/unittest/test_tvmscript_roundtrip.py | 51 ++++++++++++++++++++++
7 files changed, 121 insertions(+), 23 deletions(-)
diff --git a/src/script/printer/tir/block.cc b/src/script/printer/tir/block.cc
index 0c9289a9d2..178ed4fe75 100644
--- a/src/script/printer/tir/block.cc
+++ b/src/script/printer/tir/block.cc
@@ -180,7 +180,8 @@ Doc PrintBlock(IRDocsifier d, tir::Block block, ObjectPath
block_p, //
tir::Buffer buffer = block->alloc_buffers[i];
ObjectPath buffer_p = block_p->Attr("alloc_buffers")->ArrayIndex(i);
IdDoc lhs = DefineBuffer(buffer, *frame, d);
- ExprDoc rhs = BufferDecl(buffer, "alloc_buffer", {}, buffer_p, *frame, d);
+ ExprDoc rhs = BufferDecl(buffer, "alloc_buffer", {}, buffer_p, *frame, d,
+ BufferVarDefinition::DataPointer);
(*frame)->stmts.push_back(AssignDoc(lhs, rhs, NullOpt));
}
// Step 6. Handle `match_buffer`
diff --git a/src/script/printer/tir/buffer.cc b/src/script/printer/tir/buffer.cc
index ed8b707176..45a0dfd2ae 100644
--- a/src/script/printer/tir/buffer.cc
+++ b/src/script/printer/tir/buffer.cc
@@ -25,7 +25,7 @@ namespace script {
namespace printer {
Map<String, ExprDoc> BufferAttrs(tir::Buffer buffer, const ObjectPath&
buffer_p, const Frame& frame,
- const IRDocsifier& d) {
+ const IRDocsifier& d, BufferVarDefinition
var_definitions) {
using tvm::tir::Var;
using tvm::tir::VarNode;
Map<String, ExprDoc> kwargs;
@@ -93,11 +93,18 @@ Map<String, ExprDoc> BufferAttrs(tir::Buffer buffer, const
ObjectPath& buffer_p,
kwargs.Set("dtype", LiteralDoc::DataType(buffer->dtype,
buffer_p->Attr("dtype")));
}
// Step 3. Handle `buffer.data`
- if (!is_new_var(buffer->data)) {
+ bool is_inline_data = false;
+ if (is_new_var(buffer->data)) {
+ if (var_definitions >= BufferVarDefinition::DataPointer) {
+ is_inline_data = try_inline_def(buffer->data, buffer_p->Attr("data"),
[=]() {
+ return d->AsDoc<ExprDoc>(buffer, buffer_p)->Attr("data");
+ });
+ } else {
+ add_out_of_line_var_def(buffer->data, buffer_p->Attr("data"));
+ }
+ }
+ if (!is_inline_data) {
kwargs.Set("data", d->AsDoc<ExprDoc>(buffer->data,
buffer_p->Attr("data")));
- } else {
- try_inline_def(buffer->data, buffer_p->Attr("data"),
- [=]() { return d->AsDoc<ExprDoc>(buffer,
buffer_p)->Attr("data"); });
}
// Step 4. Handle `buffer.strides`
if (!buffer->strides.empty()) {
@@ -194,15 +201,16 @@ ExprDoc BufferCall(const ExprDoc& prefix, const
Map<String, ExprDoc>& attrs, Arr
}
ExprDoc BufferDecl(const tir::Buffer& buffer, const String& method, const
Array<ExprDoc>& args,
- const ObjectPath& p, const Frame& frame, const IRDocsifier&
d) {
+ const ObjectPath& p, const Frame& frame, const IRDocsifier&
d,
+ BufferVarDefinition var_definitions) {
return BufferCall(/*prefix=*/TIR(d, method),
- /*attrs=*/BufferAttrs(buffer, p, frame, d),
+ /*attrs=*/BufferAttrs(buffer, p, frame, d,
var_definitions),
/*args=*/args);
}
ExprDoc BufferAttn(const tir::Buffer& buffer, const ObjectPath& p, const
Frame& frame,
const IRDocsifier& d) {
- Map<String, ExprDoc> attrs = BufferAttrs(buffer, p, frame, d);
+ Map<String, ExprDoc> attrs = BufferAttrs(buffer, p, frame, d,
BufferVarDefinition::DataPointer);
ExprDoc shape = attrs.Get("shape").value();
ExprDoc dtype =
attrs.Get("dtype").value_or(LiteralDoc::DataType(buffer->dtype,
p->Attr("dtype")));
@@ -281,7 +289,8 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) //
if (!d->IsVarDefined(buffer)) {
if (Optional<Frame> opt_f = FindLowestVarDef(buffer, d)) {
ExprDoc lhs = DefineBuffer(buffer, opt_f.value(), d);
- ExprDoc rhs = BufferDecl(buffer, "Buffer", {}, p, opt_f.value(), d);
+ ExprDoc rhs = BufferDecl(buffer, "Buffer", {}, p, opt_f.value(), d,
+ BufferVarDefinition::DataPointer);
opt_f.value()->stmts.push_back(AssignDoc(lhs, rhs, NullOpt));
}
}
@@ -298,7 +307,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ExprDoc lhs = DefineBuffer(stmt->buffer, frame, d);
ExprDoc src_buffer = d->AsDoc<ExprDoc>(stmt->source,
p->Attr("source"));
ExprDoc rhs = BufferDecl(stmt->buffer, "match_buffer", {src_buffer},
p->Attr("buffer"),
- d->frames.back(), d);
+ d->frames.back(), d,
BufferVarDefinition::MatchBuffer);
return AssignDoc(lhs, rhs, NullOpt);
});
diff --git a/src/script/printer/tir/function.cc
b/src/script/printer/tir/function.cc
index a8445f23df..19cc67cb14 100644
--- a/src/script/printer/tir/function.cc
+++ b/src/script/printer/tir/function.cc
@@ -119,7 +119,8 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
ExprDoc param_doc = args[i]->lhs;
ObjectPath buffer_p = p->Attr("buffer_map")->MapValue(param);
ExprDoc lhs = DefineBuffer(buffer, *f, d);
- ExprDoc rhs = BufferDecl(buffer, "match_buffer", {param_doc},
buffer_p, *f, d);
+ ExprDoc rhs = BufferDecl(buffer, "match_buffer", {param_doc},
buffer_p, *f, d,
+ BufferVarDefinition::MatchBuffer);
(*f)->stmts.push_back(AssignDoc(lhs, rhs, NullOpt));
}
}
@@ -152,7 +153,8 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
tir::Buffer buffer = root_block->alloc_buffers[i];
ObjectPath buffer_p =
root_block_p->Attr("alloc_buffers")->ArrayIndex(i);
IdDoc lhs = DefineBuffer(buffer, *f, d);
- ExprDoc rhs = BufferDecl(buffer, "alloc_buffer", {}, buffer_p, *f,
d);
+ ExprDoc rhs = BufferDecl(buffer, "alloc_buffer", {}, buffer_p, *f, d,
+ BufferVarDefinition::DataPointer);
(*f)->stmts.push_back(AssignDoc(lhs, rhs, NullOpt));
}
AsDocBody(root_block->body, root_block_p->Attr("body"), f->get(), d);
diff --git a/src/script/printer/tir/stmt.cc b/src/script/printer/tir/stmt.cc
index 384ad6a940..01899d9001 100644
--- a/src/script/printer/tir/stmt.cc
+++ b/src/script/printer/tir/stmt.cc
@@ -126,16 +126,23 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return WhileDoc(cond, (*f)->stmts);
});
+namespace {
+Doc DeclBufferDoc(tir::DeclBuffer stmt, ObjectPath p, IRDocsifier d,
+ BufferVarDefinition var_definitions) {
+ bool concise = AllowConciseScoping(d);
+ ExprDoc rhs = BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"),
d->frames.back(), d,
+ var_definitions);
+ With<TIRFrame> f(d, stmt);
+ ExprDoc lhs = DefineBuffer(stmt->buffer, *f, d);
+ AsDocBody(stmt->body, p->Attr("body"), f->get(), d);
+ return DoConciseScoping(lhs, rhs, &(*f)->stmts, concise);
+}
+} // namespace
+
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tir::DeclBuffer>( //
"", [](tir::DeclBuffer stmt, ObjectPath p, IRDocsifier d) -> Doc {
- bool concise = AllowConciseScoping(d);
- ExprDoc rhs =
- BufferDecl(stmt->buffer, "decl_buffer", {}, p->Attr("buffer"),
d->frames.back(), d);
- With<TIRFrame> f(d, stmt);
- ExprDoc lhs = DefineBuffer(stmt->buffer, *f, d);
- AsDocBody(stmt->body, p->Attr("body"), f->get(), d);
- return DoConciseScoping(lhs, rhs, &(*f)->stmts, concise);
+ return DeclBufferDoc(stmt, p, d, BufferVarDefinition::None);
});
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
@@ -198,7 +205,8 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
"", [](tir::Allocate stmt, ObjectPath stmt_p, IRDocsifier d) -> Doc {
bool concise = AllowConciseScoping(d);
if (d->cfg->syntax_sugar && IsAllocateDeclBufferPattern(stmt.get()))
{
- return d->AsDoc(stmt->body, stmt_p->Attr("body"));
+ return DeclBufferDoc(Downcast<tir::DeclBuffer>(stmt->body),
stmt_p->Attr("body"), d,
+ BufferVarDefinition::DataPointer);
}
Array<ExprDoc> args;
Array<String> kwargs_keys;
diff --git a/src/script/printer/tir/utils.h b/src/script/printer/tir/utils.h
index cee5fbd0f0..7bdba2b1c6 100644
--- a/src/script/printer/tir/utils.h
+++ b/src/script/printer/tir/utils.h
@@ -177,6 +177,29 @@ inline std::string ReprPrintTIR(const ObjectRef& obj,
const PrinterConfig& cfg)
return Docsify(obj, d, *f, cfg);
}
+/* \brief Specify which variables are defined along with the buffer
+ *
+ * Depending on the context, defining a buffer may define additional
+ * variables associated with the buffer.
+ */
+enum class BufferVarDefinition {
+ // All parameters in the buffer must be defined prior to this call.
+ // For example, DeclBuffer.
+ None,
+
+ // The data pointer is defined along with the buffer, but buffer
+ // parameters (shape/stride/elem_offset) must be defined prior to
+ // use. For example, `BlockNode::alloc_buffers`, or the
+ // syntax-sugar representation of an `Allocate`/`DeclBuffer` pair.
+ DataPointer,
+
+ // The data pointer is defined along with the buffer, along with any
+ // buffer parameters (shape/stride/elem_offset) that have not
+ // previously been defined. For example,
+ // `BlockNode::match_buffers`, or the `PrimFuncNode::buffer_map`.
+ MatchBuffer,
+};
+
/*!
* \brief Declare and define a buffer
* \param buffer The buffer to be defined
@@ -185,10 +208,13 @@ inline std::string ReprPrintTIR(const ObjectRef& obj,
const PrinterConfig& cfg)
* \param p The object path
* \param f The frame
* \param d The IRDocsifier
+ * \param var_definitions Which variables are implicitly defined with
+ * the buffer.
* \return The ExprDoc corresponding to the buffer declaration
*/
ExprDoc BufferDecl(const tir::Buffer& buffer, const String& method, const
Array<ExprDoc>& args,
- const ObjectPath& p, const Frame& frame, const IRDocsifier&
d);
+ const ObjectPath& p, const Frame& frame, const IRDocsifier&
d,
+ BufferVarDefinition var_definitions);
/*!
* \brief Declare and define a buffer as annotation
diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py
b/tests/python/unittest/test_tvmscript_printer_tir.py
index 8427754db7..d78ba70f09 100644
--- a/tests/python/unittest/test_tvmscript_printer_tir.py
+++ b/tests/python/unittest/test_tvmscript_printer_tir.py
@@ -373,7 +373,8 @@ def test_decl_buffer():
_assert_print(
obj,
"""
-with T.decl_buffer((10, 10)) as buffer:
+v = T.handle("float32", "global")
+with T.decl_buffer((10, 10), data=v) as buffer:
T.evaluate(0)
""",
)
diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py
b/tests/python/unittest/test_tvmscript_roundtrip.py
index 7eee601358..2ea7d3ec65 100644
--- a/tests/python/unittest/test_tvmscript_roundtrip.py
+++ b/tests/python/unittest/test_tvmscript_roundtrip.py
@@ -3808,6 +3808,53 @@ def subroutine_call():
return mod
+def undefined_data_ptr_in_decl_buffer():
+ """The T.decl_buffer syntax should not introduce an Allocate
+
+ While T.decl_buffer can be used to represent an
+ Allocate/DeclBuffer pair, performing a round-trip through
+ TVMScript should not introduce an Allocate node.
+ """
+
+ @T.prim_func
+ def func():
+ data_ptr = T.handle("float32")
+ buf = T.decl_buffer(shape=[1], dtype="float32", data=data_ptr)
+ T.evaluate(buf[0])
+
+ return func
+
+
+def undefined_shape_in_decl_buffer():
+ @T.prim_func
+ def func():
+ size = T.int32()
+ buf = T.decl_buffer(shape=[size], dtype="float32")
+ T.evaluate(buf[0])
+
+ return func
+
+
+def undefined_stride_in_decl_buffer():
+ @T.prim_func
+ def func():
+ stride = T.int32()
+ buf = T.decl_buffer(shape=[1], dtype="float32", strides=[stride])
+ T.evaluate(buf[0])
+
+ return func
+
+
+def undefined_elem_offset_in_decl_buffer():
+ @T.prim_func
+ def func():
+ elem_offset = T.int32()
+ buf = T.decl_buffer(shape=[1], dtype="float32",
elem_offset=elem_offset)
+ T.evaluate(buf[0])
+
+ return func
+
+
ir_generator = tvm.testing.parameter(
launch_env_thread,
opt_gemm_normalize,
@@ -3878,6 +3925,10 @@ ir_generator = tvm.testing.parameter(
ir_module_with_attrs,
nested_seqstmt,
subroutine_call,
+ undefined_data_ptr_in_decl_buffer,
+ undefined_shape_in_decl_buffer,
+ undefined_stride_in_decl_buffer,
+ undefined_elem_offset_in_decl_buffer,
)
