pyscf 2.7.0__py3-none-macosx_11_0_arm64.whl → 2.9.0__py3-none-macosx_11_0_arm64.whl

pyscf/lib/np_helper/imatcopy.c

@@ -0,0 +1,360 @@
+/* Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
+  
+   Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+ 
+        http://www.apache.org/licenses/LICENSE-2.0
+ 
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+
+ *
+ * Author: Christopher Hillenbrand <chillenbrand15@gmail.com>
+ */
+
+#include <complex.h>
+#include <math.h>
+#include "np_helper.h"
+
+const int TILESIZE = 32;
+const int TILESIZE_CPLX = 16;
+
+/*
+ * Calculate the largest integer i such that
+ *  i * (i - 1) / 2 <= ijouter.
+ */
+static inline int uncollapse_loop_index(const long long ijouter)
+{
+  return (int) floor((sqrt(0.25 + 2.0 * ijouter) + 0.5));
+}
+
+static inline void dtranspose_scale_tile_offdiag(double *A, const int ii,
+                                                 const int jj,
+                                                 const size_t lda_w,
+                                                 const double alpha) {
+  for (int j = jj; j < jj + TILESIZE; j++) {
+#pragma omp simd
+    for(int i = ii; i < ii + TILESIZE; i++) {
+      const double tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = alpha * A[j * lda_w + i];
+      A[j * lda_w + i] = alpha * tmp;
+    }
+  }
+}
+
+static inline void dtranspose_scale_tile_diag(double *A, const int ii,
+                                              const int jj, const size_t lda_w,
+                                              const double alpha) {
+  for (int j = jj; j < jj + TILESIZE; j++) {
+#pragma omp simd
+    for(int i = ii; i < j; i++) {
+      const double tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = alpha * A[j * lda_w + i];
+      A[j * lda_w + i]= alpha * tmp;
+    }
+    A[j * lda_w + j] *= alpha;
+  }
+}
+
+static inline void dtranspose_tile_diag(double *A, const int ii, const int jj,
+                                        const size_t lda_w) {
+  for (int j = jj; j < jj + TILESIZE; j++) {
+#pragma omp simd
+    for(int i = ii; i < j; i++) {
+      const double tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = A[j * lda_w + i];
+      A[j * lda_w + i] = tmp;
+    }
+  }
+}
+
+static inline void ztranspose_scale_tile_offdiag(double complex *A,
+                                                 const int ii, const int jj,
+                                                 const size_t lda_w,
+                                                 const double complex alpha) {
+  for (int j = jj; j < jj + TILESIZE_CPLX; j++) {
+#pragma omp simd
+    for(int i = ii; i < ii + TILESIZE_CPLX; i++) {
+      const double complex tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = alpha * A[j * lda_w + i];
+      A[j * lda_w + i] = alpha * tmp;
+    }
+  }
+}
+
+static inline void ztranspose_scale_tile_diag(double complex *A, const int ii,
+                                              const int jj, const size_t lda_w,
+                                              const double complex alpha) {
+  for (int j = jj; j < jj + TILESIZE_CPLX; j++) {
+#pragma omp simd
+    for(int i = ii; i < j; i++) {
+      const double complex tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = alpha * A[j * lda_w + i];
+      A[j * lda_w + i] = alpha * tmp;
+    }
+    A[j * lda_w + j] *= alpha;
+  }
+}
+
+static inline void ztranspose_tile_diag(double complex *A, const int ii,
+                                       const int jj, const size_t lda_w) {
+  for (int j = jj; j < jj + TILESIZE_CPLX; j++) {
+#pragma omp simd
+    for(int i = ii; i < j; i++) {
+      const double complex tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = A[j * lda_w + i];
+      A[j * lda_w + i] = tmp;
+    }
+  }
+}
+
+/*
+ * In-place parallel matrix transpose, double version.
+ * See https://colfaxresearch.com/multithreaded-transposition-of-square-matrices-with-common-code-for-intel-xeon-processors-and-intel-xeon-phi-coprocessors/
+ */
+void NPomp_d_itranspose_scale(const int n, const double alpha, double *A, int lda)
+{
+  const int nclean = n - n % TILESIZE;
+  const int ntiles = nclean / TILESIZE;
+  const size_t lda_w = (size_t) lda;
+
+#pragma omp parallel
+  {
+
+/*
+ *   ---------------------------
+ *  |       ******************  |
+ *  |       ******************  |
+ *  |       ******************  |
+ *  | ******      ************  |
+ *  | ******      ************  |
+ *  | ******      ************  |
+ *  | ************      ******  |
+ *  | ************      ******  |
+ *  | ************      ******  |
+ *  | ******************        |
+ *  | ******************        |
+ *  | ******************        |
+ *  |                           |
+ *   ----------------------------
+ */
+
+/*
+ *    The following loop nest is equivalent to:
+ *    for(int iouter = 1; iouter < ntiles; iouter++)
+ *      for(int jouter = 0; jouter < iouter; jouter++)
+ *
+ *    See 10.1109/IPDPS.2017.34. 
+ */
+    int first_iteration = 1;
+    int iouter, jouter;
+#pragma omp for schedule(static) nowait
+    for(long long ijouter = 0; ijouter < (ntiles*(ntiles-1))/2; ijouter++) {
+      if(first_iteration) {
+        iouter = uncollapse_loop_index(ijouter);
+        jouter = ijouter - iouter * (iouter - 1) / 2;
+        first_iteration = 0;
+      } else {
+        jouter++;
+        if(jouter == iouter) {
+          iouter++;
+          jouter = 0;
+        }
+      }
+      dtranspose_scale_tile_offdiag(A, iouter * TILESIZE, jouter * TILESIZE, lda_w, alpha);
+    }
+
+
+/*
+ *   ---------------------------
+ *  | ******                    |
+ *  | ******                    |
+ *  | ******                    |
+ *  |       ******              |
+ *  |       ******              |
+ *  |       ******              |
+ *  |             ******        |
+ *  |             ******        |
+ *  |             ******        |
+ *  |                   ******  |
+ *  |                   ******  |
+ *  |                   ******  |
+ *  |                           |
+ *   ----------------------------
+ */
+
+    if(alpha != 1.0) {
+#pragma omp for schedule(static) nowait
+      for(int ii = 0; ii < nclean; ii+=TILESIZE) {
+        dtranspose_scale_tile_diag(A, ii, ii, lda_w, alpha);
+      }
+    } else {
+#pragma omp for schedule(static) nowait
+      for(int ii = 0; ii < nclean; ii+=TILESIZE) {
+        dtranspose_tile_diag(A, ii, ii, lda_w);
+      }
+    }
+
+
+/*
+ *   --------------------------
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  |                        ***|
+ *  | ***********************   |
+ *  | ***********************   |
+ *   ---------------------------
+ */
+
+#pragma omp for schedule(static) nowait
+    for(int j = 0; j < nclean; j++) {
+      for(int i = nclean; i < n; i++) {
+        const double tmp = A[i * lda_w + j];
+        A[i * lda_w + j] = alpha * A[j * lda_w + i];
+        A[j * lda_w + i] = alpha * tmp;
+      }
+    }
+  } // end parallel region
+
+/*
+ *   --------------------------
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                           |
+ *  |                        ***|
+ *  |                        ***|
+ *   ---------------------------
+ */
+
+  for(int j = nclean; j < n; j++) {
+    for(int i = nclean; i < j; i++) {
+      const double tmp = A[i * lda_w + j];
+      A[i * lda_w + j] = alpha * A[j * lda_w + i];
+      A[j * lda_w + i] = alpha * tmp;
+    }
+  }
+
+  if(alpha != 1.0) {
+    for(int i = nclean; i < n; i++) {
+      A[i * lda_w + i] *= alpha;
+    }
+  }
+
+}
+
+/*
+ * In-place parallel matrix transpose, double complex version.
+ * See https://colfaxresearch.com/multithreaded-transposition-of-square-matrices-with-common-code-for-intel-xeon-processors-and-intel-xeon-phi-coprocessors/
+ */
+void NPomp_z_itranspose_scale(const int n, const double complex *alphaptr, double complex *A, int lda)
+{
+  const double complex alpha = *alphaptr;
+  const int nclean = n - n % TILESIZE_CPLX;
+  const int ntiles = nclean / TILESIZE_CPLX;
+  const size_t lda_w = (size_t) lda;
+
+#pragma omp parallel
+  {
+
+/*
+ *    The following loop nest is equivalent to:
+ *    for(int iouter = 1; iouter < ntiles; iouter++)
+ *      for(int jouter = 0; jouter < iouter; jouter++)
+ *
+ *    See 10.1109/IPDPS.2017.34. 
+ */
+    int first_iteration = 1;
+    int iouter, jouter;
+#pragma omp for schedule(static) nowait
+    for(long long ijouter = 0; ijouter < (ntiles*(ntiles-1))/2; ijouter++) {
+      if(first_iteration) {
+        iouter = uncollapse_loop_index(ijouter);
+        jouter = ijouter - iouter * (iouter - 1) / 2;
+        first_iteration = 0;
+      } else {
+        jouter++;
+        if(jouter == iouter) {
+          iouter++;
+          jouter = 0;
+        }
+      }
+      ztranspose_scale_tile_offdiag(A, iouter * TILESIZE_CPLX, jouter * TILESIZE_CPLX, lda_w, alpha);
+    }
+
+    if(alpha != 1.0) {
+#pragma omp for schedule(static) nowait
+      for(int ii = 0; ii < nclean; ii+=TILESIZE_CPLX) {
+        ztranspose_scale_tile_diag(A, ii, ii, lda_w, alpha);
+      }
+    } else {
+#pragma omp for schedule(static) nowait
+      for(int ii = 0; ii < nclean; ii+=TILESIZE_CPLX) {
+        ztranspose_tile_diag(A, ii, ii, lda_w);
+      }
+    }
+
+#pragma omp for schedule(static) nowait
+    for(int j = 0; j < nclean; j++) {
+      for(int i = nclean; i < n; i++) {
+        const double complex tmp = A[i * lda_w + j];
+        A[i * lda_w + j] = alpha * A[j * lda_w + i];
+        A[j * lda_w + i] = alpha * tmp;
+      }
+    }
+
+  } // end parallel region
+
+  for(int j = nclean; j < n; j++) {
+    for(int i = nclean; i < j; i++) {
+        const double complex tmp = A[i * lda_w + j];
+        A[i * lda_w + j] = alpha * A[j * lda_w + i];
+        A[j * lda_w + i] = alpha * tmp;
+    }
+  }
+
+  if(alpha != 1.0) {
+    for(int i = nclean; i < n; i++) {
+      A[i * lda_w + i] *= alpha;
+    }
+  }
+}
+
+
+
+/*
+ * Batched versions for 3D tensors
+ */
+
+void NPomp_dtensor_itranspose_scale021(const long long matstride, int nmat, int n, const double alpha,
+                                      double *A, int lda)
+{
+  for (int imat = 0; imat < nmat; imat++) {
+    NPomp_d_itranspose_scale(n, alpha, A + imat * matstride, lda);
+  }
+}
+
+void NPomp_ztensor_itranspose_scale021(const long long matstride, int nmat, int n, const double complex *alpha,
+                                      double complex *A, int lda)
+{
+    for (int imat = 0; imat < nmat; imat++) {
+      NPomp_z_itranspose_scale(n, alpha, A + imat * matstride, lda);
+    }
+}

pyscf/lib/np_helper/np_helper.c

@@ -14,6 +14,7 @@
  */
 
 #include <stdlib.h>
+#include <complex.h>
 #include "np_helper/np_helper.h"
 
 void NPdset0(double *p, const size_t n)
@@ -47,3 +48,96 @@ void NPzcopy(double complex *out, const double complex *in, const size_t n)
                 out[i] = in[i];
         }
 }
+
+/*
+ * These are mostly useful for first-touch array allocation on NUMA systems.
+ * Use with numpy.empty.
+ */
+void NPomp_dset0(const size_t n, double *out)
+{
+#pragma omp parallel for schedule(static)
+        for (size_t i = 0; i < n; i++) {
+                out[i] = 0.0;
+        }
+}
+
+void NPomp_zset0(const size_t n, double complex *out)
+{
+#pragma omp parallel for schedule(static)
+        for (size_t i = 0; i < n; i++) {
+                out[i] = 0.0;
+        }
+}
+
+
+/*
+ * Copy a double precision matrix with multithreading.
+ */
+void NPomp_dcopy(const size_t m,
+                 const size_t n,
+                 const double *__restrict in, const size_t in_stride,
+                 double *__restrict out, const size_t out_stride)
+{
+#pragma omp parallel for schedule(static)
+        for (size_t i = 0; i < m; i++) {
+#pragma omp simd
+                for (size_t j = 0; j < n; j++) {
+                        out[i * out_stride + j] = in[i * in_stride + j];
+                }
+        }
+}
+
+/*
+ * Copy a complex double precision matrix with multithreading.
+ */
+void NPomp_zcopy(const size_t m,
+                 const size_t n,
+                 const double complex *__restrict in, const size_t in_stride,
+                 double complex *__restrict out, const size_t out_stride)
+{
+#pragma omp parallel for schedule(static)
+        for (size_t i = 0; i < m; i++) {
+#pragma omp simd
+                for (size_t j = 0; j < n; j++) {
+                        out[i * out_stride + j] = in[i * in_stride + j];
+                }
+        }
+}
+
+/*
+ * Elementwise multiplication of two double matrices.
+ * B <- A \circ B
+ */
+void NPomp_dmul(const size_t m,
+                const size_t n,
+                const double *__restrict a, const size_t a_stride,
+                double *__restrict b, const size_t b_stride,
+                double *__restrict out, const size_t out_stride)
+{
+#pragma omp parallel for schedule(static)
+        for (size_t i = 0; i < m; i++) {
+#pragma omp simd
+                for (size_t j = 0; j < n; j++) {
+                        out[i * out_stride + j] = b[i * b_stride + j] * a[i * a_stride + j];
+                }
+        }
+}
+
+/*
+ * Elementwise multiplication of two complex double matrices.
+ * B <- A \circ B
+ */
+void NPomp_zmul(const size_t m,
+                const size_t n,
+                const double complex *__restrict a, const size_t a_stride,
+                double complex *__restrict b, const size_t b_stride,
+                double complex *__restrict out, const size_t out_stride)
+{
+#pragma omp parallel for schedule(static)
+        for (size_t i = 0; i < m; i++) {
+#pragma omp simd
+                for (size_t j = 0; j < n; j++) {
+                        out[i * out_stride + j] = b[i * b_stride + j] * a[i * a_stride + j];
+                }
+        }
+}

pyscf/lib/np_helper/np_helper.h

@@ -16,6 +16,7 @@
  * Author: Qiming Sun <osirpt.sun@gmail.com>
  */
 
+#include <stdlib.h>
 #include <complex.h>
 
 #define BLOCK_DIM    104
@@ -46,6 +47,13 @@ void NPztranspose(int n, int m, double complex *a, double complex *at);
 void NPdtranspose_021(int *shape, double *a, double *at);
 void NPztranspose_021(int *shape, double complex *a, double complex *at);
 
+void NPomp_d_itranspose_scale(const int n, const double alpha, double *A, int lda);
+void NPomp_z_itranspose_scale(const int n, const double complex *alphaptr, double complex *A, int lda);
+void NPomp_dtensor_itranspose_scale021(const long long matstride, int nmat, int n, const double alpha,
+                                      double *A, int lda);
+void NPomp_ztensor_itranspose_scale021(const long long matstride, int nmat, int n, const double complex *alpha,
+                                      double complex *A, int lda);
+
 void NPdunpack_tril_2d(int count, int n, double *tril, double *mat, int hermi);
 void NPzunpack_tril_2d(int count, int n,
                        double complex *tril, double complex *mat, int hermi);
@@ -62,6 +70,24 @@ void NPzset0(double complex *p, const size_t n);
 void NPdcopy(double *out, const double *in, const size_t n);
 void NPzcopy(double complex *out, const double complex *in, const size_t n);
 
+void NPomp_dset0(const size_t n, double *out);
+void NPomp_zset0(const size_t n, double complex *out);
+
+void NPomp_dcopy(const size_t m, const size_t n,
+                 const double *in, const size_t in_stride,
+                 double *out, const size_t out_stride);
+void NPomp_zcopy(const size_t m, const size_t n,
+                 const double complex *in, const size_t in_stride,
+                 double complex *out, const size_t out_stride);
+void NPomp_dmul(const size_t m, const size_t n,
+                const double *a, const size_t a_stride,
+                double *b, const size_t b_stride,
+                double *out, const size_t out_stride);
+void NPomp_zmul(const size_t m, const size_t n,
+                const double complex *a, const size_t a_stride,
+                double complex *b, const size_t b_stride,
+                double complex *out, const size_t out_stride);
+
 void NPdgemm(const char trans_a, const char trans_b,
              const int m, const int n, const int k,
              const int lda, const int ldb, const int ldc,