Repository: sympy/sympy
Pull request: ccode: Add capability to _print_Indexed() to use user-defined strided schemes and offsets
<!-- Please give this pull request a descriptive title. Pull requests with descriptive titles are more likely to receive reviews. Describe what you changed! A title that only references an issue number is not descriptive. -->
As mentioned in #12464 , `ccode` unrolled 2D `Indexed` objects in row-major format. This PR adds support for unrolling in column-major format as well.

Example : 
```
>>> from sympy import *
>>> A = IndexedBase('A', shape=(5,3))
>>> i, j = Idx('i'), Idx('j')
>>> ccode(A[i,j])
'A[3*i + j]'
>>> A = IndexedBase('A', shape=(5,3), ordering=1)
>>> ccode(A[i,j])
'A[i + 5*j]'
```

UPDATE : Now support for strided schemes and offset have been added
<!-- If this pull request fixes an issue please indicate which issue by typing "Fixes #NNNN" below. -->
Fixes #12464 

Problem statement:
CCodePrinter._print_Indexed assumes row-major ordering
When working with BLAS and LAPACK it's quite common to work with column-major ordering.
To avoid unneccessary transposes `_print_Indexed` should also support column-major ordering.
It is probably easiest to add either an `ordering` or `strides` property to `Indexed`.

Hints from issue discussion:
Can you show me a simple example of an `Indexed` object which would use this method when translated to C code? I have gone through [this page](http://docs.sympy.org/dev/modules/tensor/indexed.html) but any `Indexed` object I'm trying out seems to lack some or the other attribute required in `_print_Indexed`.
Sure,
```
>>> A = IndexedBase('A')
>>> A = IndexedBase('A', shape=(5,3))
>>> i, j = Idx('i'), Idx('j')
>>> ccode(A[i,j])
'A[3*i + j]'
```
@bjodah I think `CCodePrinter` assumes row major ordering because accessing matrices through rows is faster than accessing them through columns.
@SatyaPrakashDwibedi no, which is faster depends on ordering. Generally speaking, iterating over the dimension with the lowest stride gives fastest access due to cache locality.

Commit message: Rectify doctest

Patch to review:
diff --git a/sympy/printing/ccode.py b/sympy/printing/ccode.py
index ee21ea62a954..ecf2905fabfa 100644
--- a/sympy/printing/ccode.py
+++ b/sympy/printing/ccode.py
@@ -170,13 +170,33 @@ def _print_Rational(self, expr):
 
     def _print_Indexed(self, expr):
         # calculate index for 1d array
-        dims = expr.shape
-        elem = S.Zero
-        offset = S.One
-        for i in reversed(range(expr.rank)):
-            elem += expr.indices[i]*offset
-            offset *= dims[i]
-        return "%s[%s]" % (self._print(expr.base.label), self._print(elem))
+        offset = expr.offset
+        strides = expr.strides
+
+        if strides is None or isinstance(strides, str):
+            dims = expr.shape
+            elem = S.Zero
+            shift = S.One
+            if strides == 'C' or strides is None:
+                traversal = reversed(range(expr.rank))
+            elif strides == 'F':
+                traversal = range(expr.rank)
+
+            for i in traversal:
+                elem += expr.indices[i]*shift
+                shift *= dims[i]
+
+            if offset is None:
+                offset = ""
+            return "%s[%s]" % (self._print(expr.base.label),
+                               self._print(elem) + str(offset))
+        elif isinstance(strides, tuple):
+            if offset is None:
+                offset = ""
+            return "%s[%s]" % (self._print(expr.base.label),
+                               str(sum([x[0]*x[1]
+                                   for x in zip(expr.indices, strides)])) +
+                               " + " + str(offset))
 
     def _print_Idx(self, expr):
         return self._print(expr.label)
diff --git a/sympy/printing/tests/test_ccode.py b/sympy/printing/tests/test_ccode.py
index 86a46ade151b..2f5dd7d9d835 100644
--- a/sympy/printing/tests/test_ccode.py
+++ b/sympy/printing/tests/test_ccode.py
@@ -249,7 +249,7 @@ def test_ccode_settings():
 def test_ccode_Indexed():
     from sympy.tensor import IndexedBase, Idx
     from sympy import symbols
-    n, m, o = symbols('n m o', integer=True)
+    s, n, m, o = symbols('s n m o', integer=True)
     i, j, k = Idx('i', n), Idx('j', m), Idx('k', o)
 
     x = IndexedBase('x')[j]
@@ -266,7 +266,16 @@ def test_ccode_Indexed():
         assert p._print_Indexed(B) == 'B[%s]' % (i*o*m+j*o+k)
         assert p._not_c == set()
 
+        A = IndexedBase('A', shape=(5,3))[i, j]
+        assert p._print_Indexed(A) == 'A[%s]' % (3*i + j)
+        A = IndexedBase('A', shape=(5,3), strides='F')[i, j]
+        assert p._print_Indexed(A) == 'A[%s]' % (i + 5*j)
 
+        A = IndexedBase('A', shape=(29,29), strides=(1, s), offset=o)[i, j]
+        assert p._print_Indexed(A) == 'A[s*j + i + o]'
+
+        A = IndexedBase('A', strides=(s, m, n), offset=o)[i, j, k]
+        assert p._print_Indexed(A) == 'A[m*j + n*k + s*i + o]'
 
 def test_ccode_Indexed_without_looking_for_contraction():
     len_y = 5
diff --git a/sympy/tensor/indexed.py b/sympy/tensor/indexed.py
index c8b958ea566e..6eaebd28c9f1 100644
--- a/sympy/tensor/indexed.py
+++ b/sympy/tensor/indexed.py
@@ -270,6 +270,14 @@ def shape(self):
                 Shape cannot be inferred from Idx with
                 undefined range: %s""" % self))
 
+    @property
+    def strides(self):
+        return self.base._strides
+
+    @property
+    def offset(self):
+        return self.base._offset
+
     @property
     def ranges(self):
         """Returns a list of tuples with lower and upper range of each index.
@@ -363,7 +371,7 @@ class IndexedBase(Expr, NotIterable):
     is_symbol = True
     is_Atom = True
 
-    def __new__(cls, label, shape=None, **kw_args):
+    def __new__(cls, label, shape=None, strides=None, offset=None, **kw_args):
         if isinstance(label, string_types):
             label = Symbol(label)
         elif isinstance(label, Symbol):
@@ -376,11 +384,18 @@ def __new__(cls, label, shape=None, **kw_args):
         elif shape is not None:
             shape = Tuple(shape)
 
+        if not(isinstance(strides, tuple) or strides in ['C', 'F'] or
+               strides is None):
+            raise TypeError(filldedent("""Stride scheme not understood.
+                             Should be tuple or string('C' or 'F' : denotes
+                             row-major or column-major respectively."""))
         if shape is not None:
             obj = Expr.__new__(cls, label, shape, **kw_args)
         else:
             obj = Expr.__new__(cls, label, **kw_args)
         obj._shape = shape
+        obj._strides = strides
+        obj._offset = offset
         return obj
 
     def __getitem__(self, indices, **kw_args):
@@ -421,6 +436,44 @@ def shape(self):
         """
         return self._shape
 
+    @property
+    def strides(self):
+        """Returns the strided scheme for the ``IndexedBase`` object.
+
+        Normally this is a tuple denoting the number of
+        steps to take in the respective dimension when traversing
+        an array. For code generation purposes strides='C' and
+        strides='F' can also be used.
+
+        strides='C' would mean that code printer would unroll
+        in row-major order and 'F' means unroll in column major
+        order.
+
+        """
+        return self._strides
+
+    @property
+    def offset(self):
+        """Returns the offset for the ``IndexedBase`` object.
+
+        This is the value added to the resulting index when the
+        2D Indexed object is unrolled to a 1D form. Used in code
+        generation.
+
+        Examples
+        ==========
+        >>> from sympy.printing import ccode
+        >>> from sympy.tensor import IndexedBase, Idx
+        >>> from sympy import symbols
+        >>> l, m, n, o = symbols('l m n o', integer=True)
+        >>> A = IndexedBase('A', strides=(l, m, n), offset=o)
+        >>> i, j, k = Idx('i'), Idx('j'), Idx('k')
+        >>> ccode(A[i, j, k])
+        'A[l*i + m*j + n*k + o]'
+
+        """
+        return self._offset
+
     @property
     def label(self):
         """Returns the label of the ``IndexedBase`` object.


Code review comments:
diff_hunk: @@ -170,13 +170,33 @@ def _print_Rational(self, expr):
 
     def _print_Indexed(self, expr):
         # calculate index for 1d array
-        dims = expr.shape
-        elem = S.Zero
-        offset = S.One
-        for i in reversed(range(expr.rank)):
-            elem += expr.indices[i]*offset
-            offset *= dims[i]
-        return "%s[%s]" % (self._print(expr.base.label), self._print(elem))
+        offset = expr.offset
+        strides = expr.strides
+
+        if strides is None or isinstance(strides, str):
+            dims = expr.shape
+            elem = S.Zero
+            shift = S.One
+            if strides == 'C' or strides is None:
+                traversal = reversed(range(expr.rank))
+            elif strides == 'F':
+                traversal = range(expr.rank)
+
+            for i in traversal:
+                elem += expr.indices[i]*shift
+                shift *= dims[i]
+
+            if offset is None:
+                offset = ""
+            return "%s[%s]" % (self._print(expr.base.label),
+                               self._print(elem) + str(offset))
+        elif isinstance(strides, tuple):
+            if offset is None:
+                offset = ""
+            return "%s[%s]" % (self._print(expr.base.label),
+                               str(sum([x[0]*x[1]
path: sympy/printing/ccode.py
text: @user1:
instead of `str` you can use `self._print` here too (it is most probably safe to use `str`, but better safe than sorry if the user calculates indices in some fancy way)

diff_hunk: @@ -170,13 +170,33 @@ def _print_Rational(self, expr):
 
     def _print_Indexed(self, expr):
         # calculate index for 1d array
-        dims = expr.shape
-        elem = S.Zero
-        offset = S.One
-        for i in reversed(range(expr.rank)):
-            elem += expr.indices[i]*offset
-            offset *= dims[i]
-        return "%s[%s]" % (self._print(expr.base.label), self._print(elem))
+        offset = expr.offset
+        strides = expr.strides
+
+        if strides is None or isinstance(strides, str):
+            dims = expr.shape
+            elem = S.Zero
+            shift = S.One
+            if strides == 'C' or strides is None:
+                traversal = reversed(range(expr.rank))
+            elif strides == 'F':
+                traversal = range(expr.rank)
+
+            for i in traversal:
+                elem += expr.indices[i]*shift
+                shift *= dims[i]
+
+            if offset is None:
+                offset = ""
+            return "%s[%s]" % (self._print(expr.base.label),
+                               self._print(elem) + str(offset))
+        elif isinstance(strides, tuple):
+            if offset is None:
+                offset = ""
path: sympy/printing/ccode.py
text: @user1:
you can let offset be 0 by default and add it to the sum, i.e. `sum(...) + offset`

diff_hunk: @@ -170,13 +170,33 @@ def _print_Rational(self, expr):
 
     def _print_Indexed(self, expr):
         # calculate index for 1d array
-        dims = expr.shape
-        elem = S.Zero
-        offset = S.One
-        for i in reversed(range(expr.rank)):
-            elem += expr.indices[i]*offset
-            offset *= dims[i]
-        return "%s[%s]" % (self._print(expr.base.label), self._print(elem))
+        offset = expr.offset
+        strides = expr.strides
+
+        if strides is None or isinstance(strides, str):
+            dims = expr.shape
+            elem = S.Zero
+            shift = S.One
+            if strides == 'C' or strides is None:
+                traversal = reversed(range(expr.rank))
+            elif strides == 'F':
+                traversal = range(expr.rank)
+
+            for i in traversal:
+                elem += expr.indices[i]*shift
+                shift *= dims[i]
+
+            if offset is None:
+                offset = ""
+            return "%s[%s]" % (self._print(expr.base.label),
path: sympy/printing/ccode.py
text: @user1:
instead of having a separate return statement we can calculate the missing `strides` here and let the method use the same logic further down as when strides is given explicitly.