fix bugs

apache · Jan 18, 2023 · c0df6f7 · c0df6f7
1 parent 874d41b
commit c0df6f7
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 11 deletions.
diff --git a/src/script/printer/tir/buffer.cc b/src/script/printer/tir/buffer.cc
@@ -129,13 +129,9 @@ ExprDoc BufferDecl(const tir::Buffer& buffer, const String& method, const Array<
 ExprDoc BufferAttn(const tir::Buffer& buffer, const ObjectPath& p, const Frame& frame,
                    const IRDocsifier& d) {
   Map<String, ExprDoc> attrs = BufferAttrs(buffer, p, frame, d);
-  Array<Doc> indices_doc;
-  for (String s : {"shape", "dtype"}) {
-    if (Optional<ExprDoc> doc = attrs.Get(s)) {
-      indices_doc.push_back(doc.value());
-    }
-  }
-  return TIR("Buffer")[indices_doc];
+  ExprDoc shape = attrs.Get("shape").value();
+  ExprDoc dtype = attrs.Get("dtype").value_or(LiteralDoc::DataType(buffer->dtype));
+  return TIR("Buffer")->Call({shape, dtype}, {}, {});
 }
 
 Array<Doc> BufferIndices(const Array<PrimExpr>& indices, const ObjectPath& p,

diff --git a/tests/python/unittest/test_tvmscript_printer_syntax_sugar.py b/tests/python/unittest/test_tvmscript_printer_syntax_sugar.py
@@ -73,7 +73,7 @@ def func(a: T.handle, b: T.handle, c: T.handle):
         C = T.match_buffer(c, [128, 128, 128], dtype="uint8")
 
     expected_output = """@T.prim_func
-def main(A: T.Buffer[(128,)], B: T.Buffer[(128, 128), "int32"], C: T.Buffer[(128, 128, 128), "uint8"]):
+def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128, 128), "int32"), C: T.Buffer((128, 128, 128), "uint8")):
     T.evaluate(0)"""
     _test(func, expected_output)
 

diff --git a/tests/python/unittest/test_tvmscript_printer_tir.py b/tests/python/unittest/test_tvmscript_printer_tir.py
@@ -57,9 +57,7 @@ def test_prim_func():
         func,
         expected="""
 @T.prim_func
-def main(a: T.handle, b: T.handle):
-    A = T.match_buffer(a, (128, 128))
-    B = T.match_buffer(b, (256, 256))
+def main(A: T.Buffer((128, 128), "float32"), B: T.Buffer((256, 256), "float32")):
     T.evaluate(0)""",
     )