Sympy: Code generation for dynamic size arrays

Question

I am using Sympy for code generation.

I would like to figure out the boundaries of what Sympy can do with respect to code generation for dynamically sized arrays.

Let's say I would like to generate a simple element-wise multiplication C-function.

void elem_mul(int m, const double *A, const double *B, double *C) {
  for(int i=0; i<m; i++){
    C[i] = A[i]*B[i];
  }
}

Using sympy.Matrix I am able to generate a loop-unrolled version of this function for a fixed m:

void elem_mul_vec3(const double *A, const double *B, double *C) {
  C[0] = A[0]*B[0];
  C[1] = A[1]*B[1];
  C[2] = A[2]*B[2];
}

But for my usecase, this is useless. I am generating code, which will be called by someone else, and I do not control the size of the vectors being input. So, I need the general thing.

I have been trying to get this working using IndexedBase et. al. but no luck. This is the closest I have gotten:

import sys
import sympy as sp
import sympy.utilities.codegen
from sympy.utilities import codegen

m = sp.symbols('m', integer=True)
i = sp.Idx('i', m)
M = sp.IndexedBase('M', shape=(m))
N = sp.IndexedBase('N', shape=(m))
K = sp.IndexedBase('K', shape=(m))
dot =sp.Eq(K[i], M[i]*N[i])

rut_obj = codegen.make_routine('dot', dot, (m, M, N, K), language='C')
rcg = codegen.CCodeGen()
rcg.dump_c([rut_obj], sys.stdout, '')

But it outputs spurious loops which I can not explain, and will not compute the correct answer:

void elem_mul(int m, double *M, double *N, double *K) {
   for (int i=0; i<m; i++){
      K[i] = 0;
   }
   for (int i=0; i<m; i++){
      for (int i=0; i<m; i++){
         K[i] = K[i] + M[i]*N[i];
      }
   }
}

Is there anyone that are able to generate a simple elementwise multiplication function like the one in the top of my post using Sympy code generation?

Any help would be greatly appreciated!

Sympy: Code generation for dynamic size arrays

0 Answers0