Index and Shape Expressions
===========================

You already know that a shape can be a tuple of symbols such as
``(n, m)`` and the elements can be accessed via indexing, e.g.
``a[i, j]``. In practice, both shapes and indices may be computed
through complex expressions. We will go through several examples in this
section.

.. code:: python

    import d2ltvm
    import numpy as np
    import tvm
    from tvm import te

Matrix Transpose
----------------

Our first example is matrix transpose ``a.T``, in which we access
``a``'s elements by columns.

.. code:: python

    n = te.var('n')
    m = te.var('m')
    A = te.placeholder((n, m), name='a')
    B = te.compute((m, n), lambda i, j: A[j, i], 'b')
    s = te.create_schedule(B.op)
    tvm.lower(s, [A, B], simple_mode=True)




.. parsed-literal::
    :class: output

    produce b {
      for (i, 0, m) {
        for (j, 0, n) {
          b[((i*stride) + (j*stride))] = a[((j*stride) + (i*stride))]
        }
      }
    }



Note that the 2-D index, e.g. ``b[i,j]`` are collapsed to the 1-D index
``b[((i*n) + j)]`` to follow the C convention.

Now verify the results.

.. code:: python

    a = np.arange(12, dtype='float32').reshape((3, 4))
    b = np.empty((4, 3), dtype='float32')
    a, b = tvm.nd.array(a), tvm.nd.array(b)
    
    mod = tvm.build(s, [A, B])
    mod(a, b)
    print(a)
    print(b)


.. parsed-literal::
    :class: output

    [[ 0.  1.  2.  3.]
     [ 4.  5.  6.  7.]
     [ 8.  9. 10. 11.]]
    [[ 0.  4.  8.]
     [ 1.  5.  9.]
     [ 2.  6. 10.]
     [ 3.  7. 11.]]


Reshaping
---------

Next let's use expressions for indexing. The following code block
reshapes a 2-D array ``a`` (``n`` by ``m`` as defined above) to 1-D
(just like ``a.reshape(-1)`` in NumPy). Note how we convert the 1-D
index ``i`` to the 2-D index ``[i//m, i%m]``.

.. code:: python

    B = te.compute((m*n, ), lambda i: A[i//m, i%m], name='b')
    s = te.create_schedule(B.op)
    tvm.lower(s, [A, B], simple_mode=True)




.. parsed-literal::
    :class: output

    produce b {
      for (i, 0, (m*n)) {
        b[i] = a[((floordiv(i, m)*stride) + (floormod(i, m)*stride))]
      }
    }



Since an :math:`n`-D array is essentially listed in the memory as a 1-D
array, the generated code does not rearrange the data sequence, but it
simplifies the index expression from 2-D (``(i//m)*m + i%m``) to 1-D
(``i``) to improve the efficiency.

We can implement a general 2-D reshape function as well.

.. code:: python

    p, q = te.var('p'), te.var('q')
    B = te.compute((p, q), lambda i, j: A[(i*q+j)//m, (i*q+j)%m], name='b')
    s = te.create_schedule(B.op)
    tvm.lower(s, [A, B], simple_mode=True)




.. parsed-literal::
    :class: output

    produce b {
      for (i, 0, p) {
        for (j, 0, q) {
          b[((i*stride) + (j*stride))] = a[((floordiv(((i*q) + j), m)*stride) + (floormod(((i*q) + j), m)*stride))]
        }
      }
    }



When testing the results, we should be aware that we put no constraint
on the output shape, which can have an arbitrary shape ``(p, q)``, and
therefore TVM will not be able to check if :math:`qp = nm` for us. For
example, in the following example we created a ``b`` with size (20)
larger than ``a`` (12), then only the first 12 elements in ``b`` are
from ``a``, others are uninitialized values.

.. code:: python

    mod = tvm.build(s, [A, B])
    a = np.arange(12, dtype='float32').reshape((3,4))
    b = np.zeros((5, 4), dtype='float32')
    a, b = tvm.nd.array(a), tvm.nd.array(b)
    
    mod(a, b)
    print(b)


.. parsed-literal::
    :class: output

    [[0.000000e+00 1.000000e+00 2.000000e+00 3.000000e+00]
     [4.000000e+00 5.000000e+00 6.000000e+00 7.000000e+00]
     [8.000000e+00 9.000000e+00 1.000000e+01 1.100000e+01]
     [8.722636e-23 3.066461e-41 9.108440e-44 0.000000e+00]
     [8.743578e-23 3.066461e-41 8.741932e-23 3.066461e-41]]


Slicing
-------

Now let's consider a special slicing operator ``a[bi::si, bj::sj]``
where ``bi``, ``bj``, ``si`` and ``sj`` can be specified later. Now the
output shape needs to be computed based on the arguments. In addition,
we need to pass the variables ``bi``, ``bj``, ``si`` and ``sj`` as
arguments when compiling the module.

.. code:: python

    bi, bj, si, sj = [te.var(name) for name in ['bi', 'bj', 'si', 'sj']]
    B = te.compute(((n-bi)//si, (m-bj)//sj), lambda i, j: A[i*si+bi, j*sj+bj], name='b')
    s = te.create_schedule(B.op)
    mod = tvm.build(s, [A, B, bi, si, bj, sj])

Now test two cases to verify the correctness.

.. code:: python

    b = tvm.nd.array(np.empty((1, 3), dtype='float32'))
    mod(a, b, 1, 2, 1, 1)
    np.testing.assert_equal(b.asnumpy(), a.asnumpy()[1::2, 1::1])
    
    b = tvm.nd.array(np.empty((1, 2), dtype='float32'))
    mod(a, b, 2, 1, 0, 2)
    np.testing.assert_equal(b.asnumpy(), a.asnumpy()[2::1, 0::2])

Summary
-------

-  Both shape dimensions and indices can be expressions with variables.
-  If a variable doesn't only appear in the shape tuple, we need to pass
   it as an argument when compiling.