Repository: apache/tvm

Issue:
[Bug] DomainTouchedAccessMap does not preserve buffer order
When using `te.extern_primfunc`, we need to ensure the tensors we pass in are the same size as the buffers we are loading them into. This is done by `python/tvm/te/operation.py`, which has [the following code](https://github.com/apache/tvm/blob/0e395c389ccd173cf6c1f254b47a81e715762626/python/tvm/te/operation.py#L364):

```python
def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
    access_map = {
        k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
    }
    in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
    input_buffers = in_buffers
    # [some lines omitted for brevity]
    for obuf in out_buffers:

    assert len(input_buffers) == len(input_tensors), (
        "The number of provided input input_tensors does not match the number of ",
        "input buffers in the primfunc",
    )
    for tensor, buffer in zip(input_tensors, input_buffers):
        assert len(tensor.shape) == len(buffer.shape)
        for d1, d2 in zip(tensor.shape, buffer.shape):
            assert d1 == d2, (
                "The input input_tensors provided do not match the input buffers in the ",
                "primfunc. Please check that the order of input te.Input_Tensors and the ",
                "order of the primfunc variables in the params list agree.",
            )
```

Specifically, by using `zip` this code requires that the order of tensors in the list `input_tensors` match the order of buffers in the **dictionary** `input_buffers`. It's easy to see how this might pose a problem.

However, the issue is slightly more nuanced. For Python 3.6 and above, [dictionaries preserve order](https://www.blog.pythonlibrary.org/2018/02/27/python-3-7-dictionaries-now-ordered). Furthermore, [there are unittests](https://github.com/apache/tvm/blob/main/tests/python/unittest/test_tir_te_extern_primfunc.py) for the code in question - why doesn't it always break?

While `extern_primfunc` probably shouldn't use dictionaries to store buffers since order is important, the issue isn't here. Instead, the problem is that `tvm.arith._ffi_api.DomainTouchedAccessMap` returns a `tvm.ir.container.Map`. When used with a small number of elements (e.g. 3), `tvm.ir.container.Map` happens to preserve order, but it does not guarantee this (which is why this bug is not detected by the unit tests). However, when a larger number (e.g. 5) elements are present in the map, an error is produced:

```
E           tvm._ffi.base.TVMError: Traceback (most recent call last):
E             22: TVMFuncCall
[some lines omitted for brevity]
E             File "/home/guberti/tvm/python/tvm/te/operation.py", line 440, in extern_primfunc
E               assert len(tensor.shape) == len(buffer.shape)
E           TVMError: AssertionError
```

### Suggested fix

As far as I can tell, `extern_primfunc` needs to be given the buffers in the same order as tensors. Thus, `DomainTouchedAccessMap` should be changed to not return a map, and instead return an list of tuples. I'm blocked by this bug, so I'm just gonna fix it.

* tir:arith

CC @areusch, @Lunderberg 



Discussion:

Resolving patch:
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -398,11 +398,12 @@ def before_split(a: T.handle, b: T.handle) -> None:
 
         C = te.extern_primfunc([A, B], func)
     """
-    access_map = {
-        k: tuple(v) for k, v in tvm.arith._ffi_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])]
+
+    # dt_access_map and primfunc.buffer_map are unordered, so use order from primfunc.params
+    dt_access_map = tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc)
+    ordered_buffers = [primfunc.buffer_map[param] for param in primfunc.params]
+    in_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][0])]
+    out_buffers = [buf for buf in ordered_buffers if len(dt_access_map[buf][1])]
     assert in_buffers, "PrimFunc has no input buffers"
     assert out_buffers, "PrimFunc has no output buffers"
 


Test patch:
diff --git a/tests/python/unittest/test_tir_te_extern_primfunc.py b/tests/python/unittest/test_tir_te_extern_primfunc.py
--- a/tests/python/unittest/test_tir_te_extern_primfunc.py
+++ b/tests/python/unittest/test_tir_te_extern_primfunc.py
@@ -21,10 +21,8 @@
 
 import tvm
 import tvm.testing
-from tvm import tir, te, TVMError
+from tvm import te
 from tvm.script import tir as T
-from tvm.arith import _ffi_api as _ffi_arith_api
-from tvm.tir.schedule import _ffi_api as _ffi_schedule_api
 
 
 # TODO(csullivan): Additional tests cases needed:
@@ -174,11 +172,11 @@ def verify_func_4(module):
 
 
 class TestPrimFuncs:
-    func, verify = tvm.testing.parameters(
-        [func_1, verify_func_1],
-        [func_2, verify_func_2],
-        [func_3, verify_func_3],
-        [func_4, verify_func_4],
+    func, params, verify = tvm.testing.parameters(
+        [func_1, ("A"), verify_func_1],
+        [func_2, ("C", "D"), verify_func_2],
+        [func_3, ("C", "A", "D", "E"), verify_func_3],
+        [func_4, ("C", "A", "D", "E"), verify_func_4],
     )
 
     def test_primfunc_call(self, func, verify):
@@ -186,11 +184,12 @@ def test_primfunc_call(self, func, verify):
         func = tvm.build(func, target=target)
         verify(func)
 
-    def test_te_extern_call(self, func, verify):
+    def test_te_extern_call(self, func, params, verify):
         ir_mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
         prim_func = ir_mod["main"]
 
-        input_tensors = create_input_tensors_for_primfunc(prim_func)
+        buf_name_map = {buf.name: buf for buf in func.buffer_map.values()}
+        input_tensors = [te.placeholder(buf_name_map[name].shape) for name in params]
         output = te.extern_primfunc(input_tensors, prim_func)
         rt_prim_func = te.create_prim_func(tensors_from_extern_op(output, prim_func))
         tvm.ir.assert_structural_equal(tvm.lower(prim_func), tvm.lower(rt_prim_func))
@@ -222,36 +221,5 @@ def tensors_from_extern_op(extern, func):
     return ordered_tensors
 
 
-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])]
-    assert in_buffers, "PrimFunc has no input buffers"
-    assert out_buffers, "PrimFunc has no output buffers"
-
-    outputs = []
-    inplace = []
-    inputs = in_buffers
-    for obuf in out_buffers:
-        if obuf in in_buffers:
-            inplace.append(obuf)
-        else:
-            outputs.append(obuf)
-
-    if not outputs:
-        iobuf = inplace.pop()
-        inputs.remove(iobuf)
-        outputs = [iobuf]
-
-    def create_tensors(input_buffers):
-        tensors = []
-        for buf in input_buffers:
-            t = te.placeholder(buf.shape, dtype=buf.dtype, name=buf.name + "_placeholder")
-            tensors.append(t)
-        return tensors
-
-    return create_tensors(inputs)
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main(sys.argv))