Date: 2026-05-01T10:02:23+00:00

Actual source code: fninvsqrt.c

      1: /*
      2:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
      3:    SLEPc - Scalable Library for Eigenvalue Problem Computations
      4:    Copyright (c) 2002-, Universitat Politecnica de Valencia, Spain

      6:    This file is part of SLEPc.
      7:    SLEPc is distributed under a 2-clause BSD license (see LICENSE).
      8:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
      9: */
     10: /*
     11:    Inverse square root function  x^(-1/2)
     12: */

     14: #include <slepc/private/fnimpl.h>
     15: #include <slepcblaslapack.h>

     17: static PetscErrorCode FNEvaluateFunction_Invsqrt(FN fn,PetscScalar x,PetscScalar *y)
     18: {
     19:   PetscFunctionBegin;
     20:   PetscCheck(x!=0.0,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Function not defined in the requested value");
     21: #if !defined(PETSC_USE_COMPLEX)
     22:   PetscCheck(x>0.0,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Function not defined in the requested value");
     23: #endif
     24:   *y = 1.0/PetscSqrtScalar(x);
     25:   PetscFunctionReturn(PETSC_SUCCESS);
     26: }

     28: static PetscErrorCode FNEvaluateDerivative_Invsqrt(FN fn,PetscScalar x,PetscScalar *y)
     29: {
     30:   PetscFunctionBegin;
     31:   PetscCheck(x!=0.0,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Derivative not defined in the requested value");
     32: #if !defined(PETSC_USE_COMPLEX)
     33:   PetscCheck(x>0.0,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Derivative not defined in the requested value");
     34: #endif
     35:   *y = -1.0/(2.0*PetscPowScalarReal(x,1.5));
     36:   PetscFunctionReturn(PETSC_SUCCESS);
     37: }

     39: static PetscErrorCode FNEvaluateFunctionMat_Invsqrt_Schur(FN fn,Mat A,Mat B)
     40: {
     41:   PetscBLASInt   n=0,ld,*ipiv,info;
     42:   PetscScalar    *Ba,*Wa;
     43:   PetscInt       m;
     44:   Mat            W;

     46:   PetscFunctionBegin;
     47:   PetscCall(FN_AllocateWorkMat(fn,A,&W));
     48:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
     49:   PetscCall(MatDenseGetArray(B,&Ba));
     50:   PetscCall(MatDenseGetArray(W,&Wa));
     51:   /* compute B = sqrtm(A) */
     52:   PetscCall(MatGetSize(A,&m,NULL));
     53:   PetscCall(PetscBLASIntCast(m,&n));
     54:   ld = n;
     55:   PetscCall(FNSqrtmSchur(fn,n,Ba,n,PETSC_FALSE));
     56:   /* compute B = A\B */
     57:   PetscCall(PetscMalloc1(ld,&ipiv));
     58:   PetscCallBLAS("LAPACKgesv",LAPACKgesv_(&n,&n,Wa,&ld,ipiv,Ba,&ld,&info));
     59:   SlepcCheckLapackInfo("gesv",info);
     60:   PetscCall(PetscLogFlops(2.0*n*n*n/3.0+2.0*n*n*n));
     61:   PetscCall(PetscFree(ipiv));
     62:   PetscCall(MatDenseRestoreArray(W,&Wa));
     63:   PetscCall(MatDenseRestoreArray(B,&Ba));
     64:   PetscCall(FN_FreeWorkMat(fn,&W));
     65:   PetscFunctionReturn(PETSC_SUCCESS);
     66: }

     68: static PetscErrorCode FNEvaluateFunctionMatVec_Invsqrt_Schur(FN fn,Mat A,Vec v)
     69: {
     70:   PetscBLASInt   n=0,ld,*ipiv,info,one=1;
     71:   PetscScalar    *Ba,*Wa;
     72:   PetscInt       m;
     73:   Mat            B,W;

     75:   PetscFunctionBegin;
     76:   PetscCall(FN_AllocateWorkMat(fn,A,&B));
     77:   PetscCall(FN_AllocateWorkMat(fn,A,&W));
     78:   PetscCall(MatDenseGetArray(B,&Ba));
     79:   PetscCall(MatDenseGetArray(W,&Wa));
     80:   /* compute B_1 = sqrtm(A)*e_1 */
     81:   PetscCall(MatGetSize(A,&m,NULL));
     82:   PetscCall(PetscBLASIntCast(m,&n));
     83:   ld = n;
     84:   PetscCall(FNSqrtmSchur(fn,n,Ba,n,PETSC_TRUE));
     85:   /* compute B_1 = A\B_1 */
     86:   PetscCall(PetscMalloc1(ld,&ipiv));
     87:   PetscCallBLAS("LAPACKgesv",LAPACKgesv_(&n,&one,Wa,&ld,ipiv,Ba,&ld,&info));
     88:   SlepcCheckLapackInfo("gesv",info);
     89:   PetscCall(PetscFree(ipiv));
     90:   PetscCall(MatDenseRestoreArray(W,&Wa));
     91:   PetscCall(MatDenseRestoreArray(B,&Ba));
     92:   PetscCall(MatGetColumnVector(B,v,0));
     93:   PetscCall(FN_FreeWorkMat(fn,&W));
     94:   PetscCall(FN_FreeWorkMat(fn,&B));
     95:   PetscFunctionReturn(PETSC_SUCCESS);
     96: }

     98: static PetscErrorCode FNEvaluateFunctionMat_Invsqrt_DBP(FN fn,Mat A,Mat B)
     99: {
    100:   PetscBLASInt   n=0;
    101:   PetscScalar    *T;
    102:   PetscInt       m;

    104:   PetscFunctionBegin;
    105:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
    106:   PetscCall(MatDenseGetArray(B,&T));
    107:   PetscCall(MatGetSize(A,&m,NULL));
    108:   PetscCall(PetscBLASIntCast(m,&n));
    109:   PetscCall(FNSqrtmDenmanBeavers(fn,n,T,n,PETSC_TRUE));
    110:   PetscCall(MatDenseRestoreArray(B,&T));
    111:   PetscFunctionReturn(PETSC_SUCCESS);
    112: }

    114: static PetscErrorCode FNEvaluateFunctionMat_Invsqrt_NS(FN fn,Mat A,Mat B)
    115: {
    116:   PetscBLASInt   n=0;
    117:   PetscScalar    *T;
    118:   PetscInt       m;

    120:   PetscFunctionBegin;
    121:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
    122:   PetscCall(MatDenseGetArray(B,&T));
    123:   PetscCall(MatGetSize(A,&m,NULL));
    124:   PetscCall(PetscBLASIntCast(m,&n));
    125:   PetscCall(FNSqrtmNewtonSchulz(fn,n,T,n,PETSC_TRUE));
    126:   PetscCall(MatDenseRestoreArray(B,&T));
    127:   PetscFunctionReturn(PETSC_SUCCESS);
    128: }

    130: static PetscErrorCode FNEvaluateFunctionMat_Invsqrt_Sadeghi(FN fn,Mat A,Mat B)
    131: {
    132:   PetscBLASInt   n=0,ld,*ipiv,info;
    133:   PetscScalar    *Ba,*Wa;
    134:   PetscInt       m;
    135:   Mat            W;

    137:   PetscFunctionBegin;
    138:   PetscCall(FN_AllocateWorkMat(fn,A,&W));
    139:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
    140:   PetscCall(MatDenseGetArray(B,&Ba));
    141:   PetscCall(MatDenseGetArray(W,&Wa));
    142:   /* compute B = sqrtm(A) */
    143:   PetscCall(MatGetSize(A,&m,NULL));
    144:   PetscCall(PetscBLASIntCast(m,&n));
    145:   ld = n;
    146:   PetscCall(FNSqrtmSadeghi(fn,n,Ba,n));
    147:   /* compute B = A\B */
    148:   PetscCall(PetscMalloc1(ld,&ipiv));
    149:   PetscCallBLAS("LAPACKgesv",LAPACKgesv_(&n,&n,Wa,&ld,ipiv,Ba,&ld,&info));
    150:   SlepcCheckLapackInfo("gesv",info);
    151:   PetscCall(PetscLogFlops(2.0*n*n*n/3.0+2.0*n*n*n));
    152:   PetscCall(PetscFree(ipiv));
    153:   PetscCall(MatDenseRestoreArray(W,&Wa));
    154:   PetscCall(MatDenseRestoreArray(B,&Ba));
    155:   PetscCall(FN_FreeWorkMat(fn,&W));
    156:   PetscFunctionReturn(PETSC_SUCCESS);
    157: }

    159: #if defined(PETSC_HAVE_CUDA)
    160: PetscErrorCode FNEvaluateFunctionMat_Invsqrt_NS_CUDA(FN fn,Mat A,Mat B)
    161: {
    162:   PetscBLASInt   n=0;
    163:   PetscScalar    *Ba;
    164:   PetscInt       m;

    166:   PetscFunctionBegin;
    167:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
    168:   PetscCall(MatDenseCUDAGetArray(B,&Ba));
    169:   PetscCall(MatGetSize(A,&m,NULL));
    170:   PetscCall(PetscBLASIntCast(m,&n));
    171:   PetscCall(FNSqrtmNewtonSchulz_CUDA(fn,n,Ba,n,PETSC_TRUE));
    172:   PetscCall(MatDenseCUDARestoreArray(B,&Ba));
    173:   PetscFunctionReturn(PETSC_SUCCESS);
    174: }

    176: #if defined(PETSC_HAVE_MAGMA)
    177: #include <slepcmagma.h>

    179: PetscErrorCode FNEvaluateFunctionMat_Invsqrt_DBP_CUDAm(FN fn,Mat A,Mat B)
    180: {
    181:   PetscBLASInt   n=0;
    182:   PetscScalar    *T;
    183:   PetscInt       m;

    185:   PetscFunctionBegin;
    186:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
    187:   PetscCall(MatDenseCUDAGetArray(B,&T));
    188:   PetscCall(MatGetSize(A,&m,NULL));
    189:   PetscCall(PetscBLASIntCast(m,&n));
    190:   PetscCall(FNSqrtmDenmanBeavers_CUDAm(fn,n,T,n,PETSC_TRUE));
    191:   PetscCall(MatDenseCUDARestoreArray(B,&T));
    192:   PetscFunctionReturn(PETSC_SUCCESS);
    193: }

    195: PetscErrorCode FNEvaluateFunctionMat_Invsqrt_Sadeghi_CUDAm(FN fn,Mat A,Mat B)
    196: {
    197:   PetscBLASInt   n=0,ld,*ipiv;
    198:   PetscScalar    *Ba,*Wa;
    199:   PetscInt       m;
    200:   Mat            W;

    202:   PetscFunctionBegin;
    203:   PetscCall(FN_AllocateWorkMat(fn,A,&W));
    204:   if (A!=B) PetscCall(MatCopy(A,B,SAME_NONZERO_PATTERN));
    205:   PetscCall(MatDenseCUDAGetArray(B,&Ba));
    206:   PetscCall(MatDenseCUDAGetArray(W,&Wa));
    207:   /* compute B = sqrtm(A) */
    208:   PetscCall(MatGetSize(A,&m,NULL));
    209:   PetscCall(PetscBLASIntCast(m,&n));
    210:   ld = n;
    211:   PetscCall(FNSqrtmSadeghi_CUDAm(fn,n,Ba,n));
    212:   /* compute B = A\B */
    213:   PetscCall(SlepcMagmaInit());
    214:   PetscCall(PetscMalloc1(ld,&ipiv));
    215:   PetscCallMAGMA(magma_xgesv_gpu,n,n,Wa,ld,ipiv,Ba,ld);
    216:   PetscCall(PetscLogFlops(2.0*n*n*n/3.0+2.0*n*n*n));
    217:   PetscCall(PetscFree(ipiv));
    218:   PetscCall(MatDenseCUDARestoreArray(W,&Wa));
    219:   PetscCall(MatDenseCUDARestoreArray(B,&Ba));
    220:   PetscCall(FN_FreeWorkMat(fn,&W));
    221:   PetscFunctionReturn(PETSC_SUCCESS);
    222: }
    223: #endif /* PETSC_HAVE_MAGMA */
    224: #endif /* PETSC_HAVE_CUDA */

    226: static PetscErrorCode FNView_Invsqrt(FN fn,PetscViewer viewer)
    227: {
    228:   PetscBool      isascii;
    229:   char           str[50];
    230:   const char     *methodname[] = {
    231:                   "Schur method for inv(A)*sqrtm(A)",
    232:                   "Denman-Beavers (product form)",
    233:                   "Newton-Schulz iteration",
    234:                   "Sadeghi iteration"
    235:   };
    236:   const int      nmeth=PETSC_STATIC_ARRAY_LENGTH(methodname);

    238:   PetscFunctionBegin;
    239:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&isascii));
    240:   if (isascii) {
    241:     if (fn->beta==(PetscScalar)1.0) {
    242:       if (fn->alpha==(PetscScalar)1.0) PetscCall(PetscViewerASCIIPrintf(viewer,"  inverse square root: x^(-1/2)\n"));
    243:       else {
    244:         PetscCall(SlepcSNPrintfScalar(str,sizeof(str),fn->alpha,PETSC_TRUE));
    245:         PetscCall(PetscViewerASCIIPrintf(viewer,"  inverse square root: (%s*x)^(-1/2)\n",str));
    246:       }
    247:     } else {
    248:       PetscCall(SlepcSNPrintfScalar(str,sizeof(str),fn->beta,PETSC_TRUE));
    249:       if (fn->alpha==(PetscScalar)1.0) PetscCall(PetscViewerASCIIPrintf(viewer,"  inverse square root: %s*x^(-1/2)\n",str));
    250:       else {
    251:         PetscCall(PetscViewerASCIIPrintf(viewer,"  inverse square root: %s",str));
    252:         PetscCall(PetscViewerASCIIUseTabs(viewer,PETSC_FALSE));
    253:         PetscCall(SlepcSNPrintfScalar(str,sizeof(str),fn->alpha,PETSC_TRUE));
    254:         PetscCall(PetscViewerASCIIPrintf(viewer,"*(%s*x)^(-1/2)\n",str));
    255:         PetscCall(PetscViewerASCIIUseTabs(viewer,PETSC_TRUE));
    256:       }
    257:     }
    258:     if (fn->method<nmeth) PetscCall(PetscViewerASCIIPrintf(viewer,"  computing matrix functions with: %s\n",methodname[fn->method]));
    259:   }
    260:   PetscFunctionReturn(PETSC_SUCCESS);
    261: }

    263: /*MC
    264:    FNINVSQRT - FNINVSQRT = "invsqrt" - The inverse square root function $f(x)=x^{-\frac{1}{2}}$.

    266:    Level: beginner

    268: .seealso: [](sec:fn), `FN`, `FNType`, `FNSetType()`
    269: M*/

    271: SLEPC_EXTERN PetscErrorCode FNCreate_Invsqrt(FN fn)
    272: {
    273:   PetscFunctionBegin;
    274:   fn->ops->evaluatefunction          = FNEvaluateFunction_Invsqrt;
    275:   fn->ops->evaluatederivative        = FNEvaluateDerivative_Invsqrt;
    276:   fn->ops->evaluatefunctionmat[0]    = FNEvaluateFunctionMat_Invsqrt_Schur;
    277:   fn->ops->evaluatefunctionmat[1]    = FNEvaluateFunctionMat_Invsqrt_DBP;
    278:   fn->ops->evaluatefunctionmat[2]    = FNEvaluateFunctionMat_Invsqrt_NS;
    279:   fn->ops->evaluatefunctionmat[3]    = FNEvaluateFunctionMat_Invsqrt_Sadeghi;
    280: #if defined(PETSC_HAVE_CUDA)
    281:   fn->ops->evaluatefunctionmatcuda[2] = FNEvaluateFunctionMat_Invsqrt_NS_CUDA;
    282: #if defined(PETSC_HAVE_MAGMA)
    283:   fn->ops->evaluatefunctionmatcuda[1] = FNEvaluateFunctionMat_Invsqrt_DBP_CUDAm;
    284:   fn->ops->evaluatefunctionmatcuda[3] = FNEvaluateFunctionMat_Invsqrt_Sadeghi_CUDAm;
    285: #endif /* PETSC_HAVE_MAGMA */
    286: #endif /* PETSC_HAVE_CUDA */
    287:   fn->ops->evaluatefunctionmatvec[0] = FNEvaluateFunctionMatVec_Invsqrt_Schur;
    288:   fn->ops->view                      = FNView_Invsqrt;
    289:   PetscFunctionReturn(PETSC_SUCCESS);
    290: }