Skip to content

Add TypedMatrixSymbol (for usage of `MatrixSymbol` in kernels)

Stephan Seitz requested to merge seitz/pystencils:matrix-symbols into master

I don't know whether this is a good idea but SymPy supports assigning MatrixSymbols. Like

>>>A = MatrixSymbols('A', 3, 3)
>>>B = MatrixSymbols('B', 3, 3)
In [12]: pystencils.Assignment(A, B)                                                                                                 
Out[12]: A := B

With this hack I can generate code like this:

#define FUNC_PREFIX static

  2 FUNC_PREFIX void kernel(float * RESTRICT _data_y, int64_t const _size_y_0, int64_t const _size_y_1, int64_t const _size_y_2, int64_t
    const _stride_y_0, int64_t const _stride_y_1, int64_t const _stride_y_2, std::function< Vector3 < double >(int, int, int) > my_fun)
  3 {
  4    for (int ctr_0 = 0; ctr_0 < _size_y_0; ctr_0 += 1)
  5    {
  6       float * RESTRICT _data_y_00 = _data_y + _stride_y_0*ctr_0;
  7       for (int ctr_1 = 0; ctr_1 < _size_y_1; ctr_1 += 1)
  8       {
  9          float * RESTRICT _data_y_00_10 = _stride_y_1*ctr_1 + _data_y_00;
 10          for (int ctr_2 = 0; ctr_2 < _size_y_2; ctr_2 += 1)
 11          {
 12             const Vector3<double> A = my_fun(ctr_0, ctr_1, ctr_2);
 13             _data_y_00_10[_stride_y_2*ctr_2] = A[0] + A[1] + A[2];
 14          }
 15       }
 16    }
 17 }
  1 #define FUNC_PREFIX static
  2 template <class Functor_T>
  3 FUNC_PREFIX void kernel(float * RESTRICT _data_y, int64_t const _size_y_0, int64_t const _size_y_1, int64_t const _size_y_2, int64_t
    const _stride_y_0, int64_t const _stride_y_1, int64_t const _stride_y_2, Functor_T my_fun)
  4 {
  5    for (int ctr_0 = 0; ctr_0 < _size_y_0; ctr_0 += 1)
  6    {
  7       float * RESTRICT _data_y_00 = _data_y + _stride_y_0*ctr_0;
  8       for (int ctr_1 = 0; ctr_1 < _size_y_1; ctr_1 += 1)
  9       {
 10          float * RESTRICT _data_y_00_10 = _stride_y_1*ctr_1 + _data_y_00;
 11          for (int ctr_2 = 0; ctr_2 < _size_y_2; ctr_2 += 1)
 12          {
 13             const Vector3<double> A = my_fun(ctr_0, ctr_1, ctr_2);
 14             _data_y_00_10[_stride_y_2*ctr_2] = A[0] + A[1] + A[2];
 15          }
 16       }
 17    }
 18 }

from

    x, y = pystencils.fields('x, y:  float32[3d]')
    from pystencils.data_types import TypedMatrixSymbol

    A = TypedMatrixSymbol('A', 3, 1, create_type('double'), 'Vector3<double>')

    my_fun_call = DynamicFunction(TypedSymbol('my_fun',
                                              'std::function< Vector3 < double >(int, int, int) >'),
                                  A.dtype,
                                  *pystencils.x_vector(3))

    assignments = pystencils.AssignmentCollection({
        A:  my_fun_call,
        y.center: A[0] + A[1] + A[2]
    })

    ast = pystencils.create_kernel(assignments)
    pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())

Merge request reports