Rosenbrock.cc

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright (c) NumeriCor GmbH
// SPDX-License-Identifier: LicenseRef-APL-1.1
 
#define NO_INCLUDE_ROSENBROCK_CU
 
/* ----------------------------------------------------------------------------
Rosenbrock-Wolfbrandt Integration Coefficients, culled from:
http://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/19960015529_1996034599.pdf
---------------------------------------------------------------------------- */
#include <string.h>
 
#include "Rosenbrock.h"
 
#if defined HAS_ROCM_MODEL
#include <hip/hip_runtime.h>
#endif
 
namespace limpet {
 
#ifdef __cplusplus
extern "C"
{
#endif  // ifdef __cplusplus
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void rbCalcKi ( float **K, float **, float *X, int* ludI,
                void (*calcDX)(float*, float*, void*),
                void *params, float h, int N, int i );
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void rbSolver ( float **, float *, float *, int N );
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void fludcmp0  ( float **, int n, int *indx, float *);
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void flubksb0  ( float **, int n, int *indx, float *);
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void fludcmp   ( float **, int, int *, float *);
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void flubksb   ( float **, int, int *, float * );
 
 
#if defined __CUDA__ || defined __HIP__
__device__
#endif
void rbStepX ( float *X, void (*calcDX)(float*, float*, void*),
                void (*calcJ)(float**, float*, void*, int ),
                void *params, float h, int N )
{
  float  kbuf[RS_MAX_N*RS_ORDER];
  float *k[RS_ORDER];
  float  abuf[RS_MAX_N*RS_MAX_N];
  float *a[RS_MAX_N];
 
  //const float A_rs[RS_ORDER] = { 0., 0.5, 0.5, 1. };
  //const float B_rs[RS_ORDER] = { 0.25, 0.25, 0.25, 0.25 };
  const float Chi[RS_ORDER]  = { 14/3., 20/3., 4/3., 2/3. };
 
  if (N > RS_MAX_N) {
#if !defined __CUDA__ && !defined __HIP__
    printf( "Increase limit for number of equations: currently %d\n", RS_MAX_N);
#elif defined __CUDA__
    printf( "Error: Increase limit for number of equations (RS_MAX_N macro)\n");
#endif
    return;
  }
 
  memset(kbuf, 0, sizeof(kbuf));
  memset(abuf, 0, sizeof(abuf));
 
  for( int i = 0; i < RS_ORDER || i < N; i++ ) {
    if( i < RS_ORDER ) k[i] = &(kbuf[i*N]);
    if( i < N )     a[i] = &(abuf[i*N]);
  }
  calcJ(a, X, params, N);
 
  // a = 4i - hj;
  for( int i = 0; i < N; i++ )
    for( int j = 0; j < N; j++ )
      a[i][j] = ((i == j) ? 4 : 0) - h * a[i][j];
 
  int    ludI[RS_MAX_N];
  float  ludD;
  fludcmp0( a, N, ludI, &ludD );
 
  for( int ki = 0; ki < RS_ORDER; ki++)
    rbCalcKi ( k, a, X, ludI, calcDX, params, h, N, ki );
 
  for( int Ki = 0; Ki < RS_ORDER; Ki++)
    for( int row = 0; row < N; row++ )
      X[row] += h * Chi[Ki] * k[Ki][row];
 
}
 
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void rbCalcKi ( float **K, float **A, float *X, int* ludI,
                void (*calcDX)(float*, float*, void*),
                void *params, float h, int N, int i        )
{
 
  float xx[RS_MAX_N];
  const float Alpha[RS_ORDER][RS_ORDER] = {
      { 0,  0,  0,  0 },
      { 2,  0,  0,  0 },
      { 2,  2,  0,  0 },
      { 6, 10,  4,  0 } };
  const float Beta[RS_ORDER][RS_ORDER]  = {
      {  0,  0,  0,  0 },
      { -4,  0,  0,  0 },
      { -6, -10, 0,  0 },
      { -4, -12, 0,  0 } };
 
  memcpy( xx, X, N * sizeof(float));
 
  for( int j = 0; j < i; j++ ) {
    for( int row = 0; row < N; row++ ) {
      xx[row]   += h * Alpha[i][j] * K[j][row];
      K[i][row] +=      Beta[i][j] * K[j][row];
    }
  }
 
  float DXi[RS_MAX_N];
  calcDX( DXi, xx, params );
 
  for( int row = 0; row < N; row++ ) K[i][row] += DXi[row];
 
  //    rbSolver( &A[0][0], &K[RM(N,i,0)], bb, N );
  flubksb0( A, N, ludI, K[i] );
  return;
}
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void rbSolver( float **A, float *x, float *b, int N ) {
  // Purpose:    solve the system Ax = b by Gaussian elimination
 
  // Step #1: convert the system to upper-diagonal form
  for( int i = 1; i < N; i++ ) {
    for( int j = 0; j < i; j++ ) {
      // Set A[i][j] to zero.
      float K = -A[i][j] / A[j][j];
      for( int k = j; k < N; k++ ) A[i][k] += K * A[j][k];
      b[i] += K * b[j];
    }
  }
 
  // Step #2: solve for x by back-substitution
  for(int i = N - 1; i >= 0; i--) {
    x[i] = b[i];
    for(int j = i + 1; j < N; j++) x[i] -= A[i][j] * x[j];
    x[i] /= A[i][i];
  }
}
 
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void fludcmp0( float **a, int n,int *indx, float *d)/* matrices start at 0 */
{
  float *ptr[RS_MAX_N+1];
  int i;
 
  for( i=0; i<n; i++ )
    ptr[i+1] = a[i]-1;
  fludcmp( ptr, n, indx-1, d );
 
  for(i=0; i<n; i++)
    a[i] = ptr[i+1] + 1;
}
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void flubksb0( float **a, int n, int *indx, float *b)
{
  float *ptr[RS_MAX_N+1];
  int i;
 
  for( i=0; i<n; i++ )
    ptr[i+1] = a[i]-1;
 
  flubksb( ptr, n, indx-1, b-1 );
}
 
#define TINY 1.0e-20;
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void fludcmp( float **a, int n,int *indx, float *d)
{
  int i,j;
  int imax = 0,k;
  float big,dum,temp;
  float sum;
  float *dumpoint, vv[RS_MAX_N+1];
 
 
  *d=1.0;
  i = 1;
  do {
    big=0.0;
    for (j=1;j<=n;j++)
      if ((temp=fabs(a[i][j])) > big) big=temp;
    if (big == 0.0) {
      #ifndef __HIP__
      printf("Singular matrix in routine LUDCMP");
      #endif
      //exit(1);
      //
    }
    vv[i]=1.0/big;
  } while ( ++i <= n );
 
  j = 1;
  do {
    if( j > 1 ) {
      i = 1;
      do {
        sum=a[i][j];
        if ( i != 1 ) {
          k = 1;
          do {
            sum -= a[i][k]*a[k][j];
          } while ( ++k < i );
        }
        a[i][j]=sum;
      } while ( ++i < j );
    }
    big=0.0;
    i = j;
    do {
      sum=a[i][j];
      if ( j > 1 ) {
        k =1 ;
        do {
          sum -= a[i][k]*a[k][j];
        } while( ++k < j );
      }
      a[i][j]=sum;
      if ( (dum=vv[i]*fabs(sum)) >= big) {
        big=dum;
        imax=i;
      }
    } while( ++i <= n );
    if (j != imax) {    /* exchange rows */
      dumpoint=a[imax];
      a[imax]=a[j];
      a[j]=dumpoint;
      *d = -(*d);
      vv[imax]=vv[j];
    }
    indx[j]=imax;
    if (a[j][j] == 0.0) a[j][j]=TINY;
    if (j != n) {
      dum=1.0/(a[j][j]);
      for (i=j+1;i<=n;i++) a[i][j] *= dum;
    }
  } while( ++j <= n );
}
 
#undef TINY
 
#if defined MLIR_CODEGEN && (defined __CUDA__ || defined __HIP__)
__device__
#endif
void flubksb( float **a, int n, int *indx, float *b )
{
  int i,ii=0,ip;
  int j;
  float sum;
 
  for (i=1;i<=n;i++) {
    ip=indx[i];
    sum=b[ip];
    b[ip]=b[i];
    if (ii)
      for (j=ii;j<=i-1;j++)
        sum -= a[i][j]*b[j];
    else if (sum)
      ii=i;
    b[i]=sum;
  }
  for (i=n;i>=1;i--) {
    sum=b[i];
    for (j=i+1;j<=n;j++)
      sum -= a[i][j]*b[j];
    b[i]=sum/a[i][i];
  }
}
 
#ifdef __cplusplus
}
#endif  // ifdef __cplusplus
 
#if defined MLIR_CODEGEN && !defined __CUDA__ && !defined __HIP__
 
template<size_t vector_size>
void rbCalcKi(float **K, float **A, float *X, int* ludI,
              void (*calcDX)(float*, float*, void*),
              void *params, float h, int N, int i)
{
 
  float xx[RS_MAX_N * vector_size];
  const float Alpha[RS_ORDER][RS_ORDER] = {
      { 0,  0,  0,  0 },
      { 2,  0,  0,  0 },
      { 2,  2,  0,  0 },
      { 6, 10,  4,  0 } };
  const float Beta[RS_ORDER][RS_ORDER]  = {
      {  0,  0,  0,  0 },
      { -4,  0,  0,  0 },
      { -6, -10, 0,  0 },
      { -4, -12, 0,  0 } };
 
  memcpy(xx, X, N*vector_size*sizeof(float));
 
  for (int v = 0; v < vector_size; ++v) {
    for( int j = 0; j < i; j++ ) {
      for( int row = 0; row < N; row++ ) {
        xx[v*N + row] += h * Alpha[i][j] * K[v*RS_ORDER + j][row];
        K[v*RS_ORDER + i][row] += Beta[i][j] * K[v*RS_ORDER + j][row];
      }
    }
  }
 
  float DXi[RS_MAX_N * vector_size];
  calcDX( DXi, xx, params );
  for (int v = 0; v < vector_size; ++v)
    for( int row = 0; row < N; row++ )
      K[v*RS_ORDER + i][row] += DXi[v*N + row];
 
  for (int v = 0; v < vector_size; ++v)
    flubksb0(&A[v*N], N, ludI+v*N, K[v*RS_ORDER + i]);
}
 
template<size_t vector_size>
void rbStepX ( float *X, void (*calcDX)(float*, float*, void*),
               void (*calcJ)(float*, float*, void*, int ),
               void *params, float h, int N )
{
  float  kbuf[RS_MAX_N*RS_ORDER*vector_size];
  float *k[RS_ORDER*vector_size];
  float  abuf[RS_MAX_N*RS_MAX_N*vector_size];
  float *a[RS_MAX_N*vector_size];
 
  const float Chi[RS_ORDER]  = { 14/3., 20/3., 4/3., 2/3. };
 
  if(N>RS_MAX_N) {
    printf( "Increase limit for number of equations: currently %d\n", RS_MAX_N );
    return;
  }
 
  memset(kbuf, 0, sizeof(kbuf));
  memset(abuf, 0, sizeof(abuf));
 
  int K_ORDER = N > RS_ORDER ? N : RS_ORDER;
  for (int i = 0; i < vector_size; ++i) {
    for( int j = 0; j < RS_ORDER || j < N; j++ ) {
      if( j < RS_ORDER ) k[i*RS_ORDER + j] = &(kbuf[i*K_ORDER*K_ORDER + j*K_ORDER]);
      if( j < N )     a[i*N + j] = &(abuf[i*N*N + j*N]);
    }
  }
  calcJ(abuf, X, params, N);
 
  for (int i = 0; i < vector_size; ++i)
    for( int j = 0; j < N; j++ )
      for( int k = 0; k < N; k++ )
        a[i*N + j][k] = ((j == k) ? 4 : 0) - h * a[i*N + j][k];
 
  int ludI[RS_MAX_N * vector_size];
  for (int i = 0; i < vector_size; ++i) {
    float  ludD;
    fludcmp0( &a[i*N], N, ludI + i * N, &ludD );
  }
 
  for( int ki = 0; ki < RS_ORDER; ki++)
    rbCalcKi<vector_size>( k, a, X, ludI, calcDX, params, h, N, ki);
 
  for (int i = 0; i < vector_size; ++i)
    for( int Ki = 0; Ki < RS_ORDER; Ki++)
      for( int row = 0; row < N; row++ )
        X[i*N + row] += h * Chi[Ki] * k[i * RS_ORDER + Ki][row];
}
 
#ifdef __cplusplus
extern "C"
{
#endif  // ifdef __cplusplus
 
void rbStepX_2 ( float *X, void (*calcDX)(float*, float*, void*),
               void (*calcJ)(float*, float*, void*, int ),
               void *params, float h, int N )
{
 rbStepX<2>(X, calcDX, calcJ, params, h, N);
}
 
void rbStepX_4 ( float *X, void (*calcDX)(float*, float*, void*),
               void (*calcJ)(float*, float*, void*, int ),
               void *params, float h, int N )
{
 rbStepX<4>(X, calcDX, calcJ, params, h, N);
}
 
void rbStepX_8 ( float *X, void (*calcDX)(float*, float*, void*),
               void (*calcJ)(float*, float*, void*, int ),
               void *params, float h, int N )
{
 rbStepX<8>(X, calcDX, calcJ, params, h, N);
}
 
void rbStepX_16 ( float *X, void (*calcDX)(float*, float*, void*),
               void (*calcJ)(float*, float*, void*, int ),
               void *params, float h, int N )
{
 rbStepX<16>(X, calcDX, calcJ, params, h, N);
}
 
void rbStepX_32 ( float *X, void (*calcDX)(float*, float*, void*),
               void (*calcJ)(float*, float*, void*, int ),
               void *params, float h, int N )
{
 rbStepX<32>(X, calcDX, calcJ, params, h, N);
}
 
#ifdef __cplusplus
}
#endif  // ifdef __cplusplus
 
#endif
 
}  // namespace limpet