[SparseTIR] Parser, Printer, Roundtrip

[MasterJH5574]

This PR implements TVMScript parser and printer for data structures like SparseBlock, SparseBufferLoad/Store, etc.

After this PR, TVMScript roundtrip for SparseTIR is achieved.

cc @yzh119

---

To do:

• Add more roundtrip unit tests

[yzh119]

Would you mind showing a demo of the printed script in the thread?

[yzh119]

What does this attribute mean?

[MasterJH5574]

Suppose we have a sparse axis I. As you know, when defining a SparseBuffer we provide the list of axes of the buffer. And sometimes there is a T.to_dense(I) among those axis. The following is an example:

tvm/tests/python/unittest/test_tir_sparse_script_roundtrip.py

Line 33 in fb2c49f

B = T.match_sparse_buffer(b, (T.to_dense(J), K), m * k, "float32")

The fact is, without knowing which sparse axis was a dense-fixed axis created from, we can never print out T.to_dense(...) from a PrimFunc. And instead the printer will print a new T.dense_fixed(...) for that axis. This behavior causes inconsistency between the original script and the newly printed script.

Therefore, for the print purpose, the from_sparse field is introduced. It might seem an unnecessary stuff. And I haven't tested whether removing it will influence the roundtrip property (may or may not). Anyway, it doesn't hurt, and this is why I added this.

[yzh119]

So the major purpose is for printing.
Yeah another idea I can come up with is to store a relation graph among these axes.

[MasterJH5574]

It's doable. Just wonder whether the graph has other usage? IIUC there is only one kind of relation indicated by to_dense(...) (is this true?), and the graph might be unnecessary for us 🤔.

[yzh119]

Haha I'm not sure, yeah I agree that it's not necessary at this stage.

[yzh119]

Could you briefly explain how does the MarkGraphNode work?

[MasterJH5574]

When MarkGraphNode is on, the pair of Objects to be compared will be marked as graph_equal = true:

tvm/src/node/structural_equal.cc

Lines 92 to 96 in fb2c49f

	void MarkGraphNode() final {
	// need to push to pending tasks in this case
	ICHECK(!allow_push_to_stack_ && !task_stack_.empty());
	task_stack_.back().graph_equal = true;
	}

In this situation, when doing structural equality check, the checker will put the equivalent pair into two internal map containers:

tvm/src/node/structural_equal.cc

Lines 150 to 153 in fb2c49f

	if (entry.graph_equal) {
	equal_map_lhs_[entry.lhs] = entry.rhs;
	equal_map_rhs_[entry.rhs] = entry.lhs;
	}

And the mappings will provide equivalence information when we do structural equality check over two Map, even when their items have different relative orders:

tvm/src/node/structural_hash.cc

Lines 463 to 476 in fb2c49f

	static bool SEqualReduceForOMap(const MapNode* lhs, const MapNode* rhs, SEqualReducer equal) {
	for (const auto& kv : *lhs) {
	// Only allow equal checking if the keys are already mapped
	// This resolves common use cases where we want to store
	// Map<Var, Value> where Var is defined in the function
	// parameters.
	ObjectRef rhs_key = equal->MapLhsToRhs(kv.first);
	if (!rhs_key.defined()) return false;
	auto it = rhs->find(rhs_key);
	if (it == rhs->end()) return false;
	if (!equal(kv.second, it->second)) return false;
	}
	return true;
	}

The method MapLhsToRhs(...) simply looks up the mappings:

tvm/src/node/structural_equal.cc

Lines 98 to 103 in fb2c49f

	ObjectRef MapLhsToRhs(const ObjectRef& lhs) final {
	auto it = equal_map_lhs_.find(lhs);
	if (it != equal_map_lhs_.end()) return it->second;
	return ObjectRef(nullptr);
	}

---

Finally, we can safely do structural equality check among mappings sp_struct2param_map without worrying about the relative orders of the mapping items 😉.

[yzh119]

That make sense to me, btw why don't we maintain a union-find set (I mean, in TVM) for equivalence check?

[MasterJH5574]

In my understanding structural equality check only compares two objects once at a time. I think union-find set will be more helpful to represent equivalence in among many objects.

I'm not the designer of this check, though. But if I was asked to write such a check, I'd also like to keep the check as lightweight as possible (I mean, keep little information as possible).

[MasterJH5574]

Would you mind showing a demo of the printed script in the thread?

@yzh119 Sure! It's here https://www.diffchecker.com/A13klPkg.

On the left side are the original functions in the test, and on the right side are the printed functions.

[yzh119]

Nice job @MasterJH5574 , let's move on to next milestone.