Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[IR] Refactor IR Printer: HLOExpr to Doc #198

Merged
merged 3 commits into from
Mar 17, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions include/matxscript/ir/printer/doc.h
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,19 @@ class LiteralDoc : public ExprDoc {
return LiteralDoc::Str(dtype, p);
}

/*!
* \brief Create a LiteralDoc to represent string.
* \param v The string value.
* \param p The object path
*/
static LiteralDoc HLOType(const Type& v, const Optional<ObjectPath>& p) {
if (auto const* pt = v.as<PrimTypeNode>()) {
return LiteralDoc::DataType(pt->dtype, p);
}
StringRef dtype = v->GetPythonTypeName().encode();
return LiteralDoc::Str(dtype, p);
}

MATXSCRIPT_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
};

Expand Down
16 changes: 16 additions & 0 deletions src/ir/_base/cow_array_ref.cc
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
#include <matxscript/ir/_base/repr_printer.h>
#include <matxscript/ir/_base/structural_equal.h>
#include <matxscript/ir/_base/structural_hash.h>
#include <matxscript/ir/printer/doc.h>
#include <matxscript/ir/printer/ir_docsifier.h>
#include <matxscript/runtime/functor.h>
#include <matxscript/runtime/registry.h>

Expand Down Expand Up @@ -197,5 +199,19 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
p->stream << ']';
});

using namespace ::matxscript::ir::printer;
MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<Array<ObjectRef>>( //
"",
[](Array<ObjectRef> array, ObjectPath p, IRDocsifier d) -> Doc {
int n = array.size();
Array<ExprDoc> results;
results.reserve(n);
for (int i = 0; i < n; ++i) {
results.push_back(d->AsDoc<ExprDoc>(array[i], p->ArrayIndex(i)));
}
return ListDoc(results);
});

} // namespace ir
} // namespace matxscript
38 changes: 38 additions & 0 deletions src/ir/_base/cow_map_ref.cc
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@
#include <matxscript/ir/_base/repr_printer.h>
#include <matxscript/ir/_base/structural_equal.h>
#include <matxscript/ir/_base/structural_hash.h>
#include <matxscript/ir/printer/doc.h>
#include <matxscript/ir/printer/ir_docsifier.h>
#include <matxscript/runtime/functor.h>
#include <matxscript/runtime/registry.h>

Expand Down Expand Up @@ -364,5 +366,41 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(ReprPrinter, vtable)

MATX_DLL constexpr uint64_t DenseMapNode::kNextProbeLocation[];

using namespace ::matxscript::ir::printer;
MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<Map<ObjectRef, ObjectRef>>( //
"",
[](Map<ObjectRef, ObjectRef> dict, ObjectPath p, IRDocsifier d) -> Doc {
using POO = std::pair<ObjectRef, ObjectRef>;
std::vector<POO> items{dict.begin(), dict.end()};
bool is_str_map = true;
for (const auto& kv : items) {
if (!kv.first.as<StringNode>()) {
is_str_map = false;
break;
}
}
if (is_str_map) {
std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
return runtime::Downcast<StringRef>(lhs.first) <
runtime::Downcast<StringRef>(rhs.first);
});
} else {
std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
return lhs.first.get() < rhs.first.get();
});
}
int n = dict.size();
Array<ExprDoc> ks;
Array<ExprDoc> vs;
ks.reserve(n);
vs.reserve(n);
for (int i = 0; i < n; ++i) {
ks.push_back(d->AsDoc<ExprDoc>(items[i].first, p->MissingMapEntry()));
vs.push_back(d->AsDoc<ExprDoc>(items[i].second, p->MapValue(items[i].first)));
}
return DictDoc(ks, vs);
});

} // namespace ir
} // namespace matxscript
9 changes: 9 additions & 0 deletions src/ir/_base/string_ref.cc
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
#include <initializer_list>

#include <matxscript/ir/_base/reflection.h>
#include <matxscript/ir/printer/doc.h>
#include <matxscript/ir/printer/ir_docsifier.h>
#include <matxscript/runtime/container/string_helper.h>
#include <matxscript/runtime/object.h>
#include <matxscript/runtime/registry.h>
Expand Down Expand Up @@ -262,5 +264,12 @@ typename StringRef::const_reverse_iterator StringRef::rend() const {
return const_reverse_iterator(begin());
}

using namespace ::matxscript::ir::printer;
MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<StringRef>("", [](StringRef s, ObjectPath p, IRDocsifier d) -> Doc {
// TODO: optimize MultipleLines
return LiteralDoc::Str(s, p);
});

} // namespace ir
} // namespace matxscript
8 changes: 8 additions & 0 deletions src/ir/adt.cc
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,16 @@
*/
#include <matxscript/ir/adt.h>

#include <matxscript/ir/printer/doc.h>
#include <matxscript/ir/printer/ir_docsifier.h>
#include <matxscript/ir/type.h>
#include <matxscript/runtime/registry.h>

namespace matxscript {
namespace ir {

using namespace runtime;
using namespace ::matxscript::ir::printer;

Constructor::Constructor(Type ret_type,
StringRef name_hint,
Expand Down Expand Up @@ -66,6 +69,11 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
<< node->belong_to << ")";
});

MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable) //
.set_dispatch<Constructor>("", [](Constructor node, ObjectPath p, IRDocsifier d) -> Doc {
return IdDoc(node->name_hint);
});

// ClassType
ClassType::ClassType(uint64_t py_type_id,
GlobalTypeVar header,
Expand Down
Loading