[TIR] Fix extern_primfunc buffer order bug

[guberti]

This fixes #13330, which was blocking my work to write TIR schedules for microTVM.

I originally thought I'd have to change the function signature of DomainTouchedAccessMap, but I couldn't think of a way to do that cleanly. Instead, I changed extern_primfunc to use primfunc.params to create the buffer lists in the right order.

#13330 should have been caught by test_tir_te_extern_primfunc.py, but one of that test's helper functions had the same bug as extern_primfunc. I've thus modified test_tir_te_extern_primfunc.py to instantiate the input tensors a different way, allowing it to catch regressions of this issue.

Would love a look from @csullivan!

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao _{See #10317 for details}
• Built docs for commit 88c37a2 can be found here.

_{Generated by tvm-bot}

[guberti]

I'd love to add a unit test to prevent regressions - ideally in tests/python/unittest/test_tir_te_extern_primfunc.py. However, when running on main I can't get that unit test to pass. I get:

Turns out this was caused by using Python 3.8 instead of 3.7. See #13360.

[guberti]

Looks like the unit tests are failing because the create_input_tensors_for_primfunc helper function in test_tir_te_extern_primfunc.py uses the same DomainTouchedAccessMap function to order its tensors, so the orders happened to match before.

[csullivan]

Thanks @guberti! Yes, I think you'll want to make the same change in the test as well

diff --git a/tests/python/unittest/test_tir_te_extern_primfunc.py b/tests/python/unittest/test_tir_te_extern_primfunc.py
index 2675214562..a3c6645be7 100644
--- a/tests/python/unittest/test_tir_te_extern_primfunc.py
+++ b/tests/python/unittest/test_tir_te_extern_primfunc.py
@@ -224,8 +224,9 @@ def tensors_from_extern_op(extern, func):

 def create_input_tensors_for_primfunc(primfunc):
     access_map = {k: tuple(v) for k, v in _ffi_arith_api.DomainTouchedAccessMap(primfunc).items()}
-    in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
-    out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
+    ordered_buffers = [primfunc.buffer_map[param] for param in primfunc.params]
+    in_buffers = [buf for buf in ordered_buffers if len(access_map[buf][0])]
+    out_buffers = [buf for buf in ordered_buffers if len(access_map[buf][1])]
     assert in_buffers, "PrimFunc has no input buffers"
     assert out_buffers, "PrimFunc has no output buffers"

Please also add a failing (on HEAD) test that reproduces the issue so we can prevent the regression in the future. I admit the need for more tests in the header of that test. Unfortunately I have not had time to return to it recently, but I am happy to hear you are using it and appreciate the contribution!

[guberti]

@csullivan I fixed test_tir_te_extern_primfunc.py so that it fails on HEAD (and will catch regressions of this issue). I would have done that sooner, but I was blocked for a bit by #13360. Should be good to merge now!

[junrushao]

I'm happy to merge it in once it's confirmed that the TVMScript issue has been fixed.

[guberti]

@csullivan would you mind taking another look and approving if this looks good?

@junrushao I've confirmed the TVMScript issue (#13360) is fixed, can we merge this?

[junrushao]

Thanks @guberti for confirming :-) Merging it in

[junrushao]

LGTM!

[Lunderberg]

LGTM, too! I like the improvement of using primfunc's parameter list for making the tensor to buffer mapping, since the parameter list is what actually defines the order of parameters.