Fix GPU spectral bugs on AMD/ROCm (alignment, symbol visibility, adjoint kernel miscompilation)

examples/hmctest/adjointtest4.cpp

@@ -771,16 +771,10 @@ static int runTest17(int gpuDevice) {
             { 2, 7, 2, 0 },
             { 3, 6, 3, 0 }
         };
-        /* GPU writes kPaddedStateCount² values with stride kPaddedStateCount;
-         * CPU writes kStateCount² values with stride kStateCount.
-         * Allocate the padded size to prevent overflow, then re-index to S×S. */
-        int Sp = useGpu ? ((S <= 4) ? 4 : (S <= 16) ? 16 : (S <= 32) ? 32 : S) : S;
-        std::vector<double> gradBuf(Sp * Sp, 0.0);
-        beagleCalculateAdjointDerivative(inst, branchOps, 0, 0, 4, 0, 4, gradBuf.data(), NULL);
+        /* Both CPU and GPU write kStateCount² values with stride kStateCount
+         * (GPU downsamples from its internal kPaddedStateCount² buffer). */
         std::vector<double> grad(S * S, 0.0);
-        for (int ls = 0; ls < S; ls++)
-            for (int rs = 0; rs < S; rs++)
-                grad[ls*S+rs] = gradBuf[ls*Sp+rs];
+        beagleCalculateAdjointDerivative(inst, branchOps, 0, 0, 4, 0, 4, grad.data(), NULL);
 
         const char* preLabels[] = { "pre6", "pre7", "pre8", "pre9" };
         for (int b = 0; b < 4; b++) {
@@ -794,12 +788,12 @@ static int runTest17(int gpuDevice) {
         const char* bNames[4] = { "chimp", "human", "gorilla", "internal" };
         fprintf(stdout, "  Per-branch gradient G[ls,rs] (first 4x4 sub-block):\n");
         for (int b = 0; b < 4; b++) {
-            std::vector<double> gbuf(Sp * Sp, 0.0);
+            std::vector<double> gbuf(S * S, 0.0);
             beagleCalculateAdjointDerivative(inst, &branchOps[b], 0, 0, 4, 0, 1, gbuf.data(), NULL);
             fprintf(stdout, "    Branch %s (transMatrix=%d):\n", bNames[b], branchOps[b].branchTransitionMatrix);
             for (int ls = 0; ls < 4; ls++) {
                 for (int rs = 0; rs < 4; rs++)
-                    fprintf(stdout, " %12.4f", gbuf[ls*Sp+rs]);
+                    fprintf(stdout, " %12.4f", gbuf[ls*S+rs]);
                 fprintf(stdout, "\n");
             }
         }

libhmsbeagle/CPU/BeagleCPUPlugin.cpp

@@ -50,7 +50,7 @@ Plugin("CPU", "CPU")
 
 
 extern "C" {
-void* plugin_init(void){
+__attribute__((visibility("default"))) void* plugin_init(void){
 	return new beagle::cpu::BeagleCPUPlugin();
 }
 }

libhmsbeagle/CPU/BeagleCPUSSEPlugin.cpp

@@ -176,7 +176,7 @@ bool check_sse2(){
 #endif
 
 
-void* plugin_init(void){
+__attribute__((visibility("default"))) void* plugin_init(void){
 	if(!check_sse2()){
 		return NULL;	// no SSE no plugin?!
 	}

libhmsbeagle/CPU/BeagleCPUSpectralPlugin.cpp

@@ -48,7 +48,7 @@ Plugin("CPU-Spectral", "CPU-Sprectral")
 
 
 extern "C" {
-void* plugin_init(void){
+__attribute__((visibility("default"))) void* plugin_init(void){
 	return new beagle::cpu::BeagleCPUSpectralPlugin();
 }
 }

libhmsbeagle/CPU/CMakeLists.txt

@@ -25,6 +25,8 @@ SET_TARGET_PROPERTIES(hmsbeagle-cpu
 	SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
 
+target_compile_options(hmsbeagle-cpu PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)
+
 if(BUILD_SSE)
 add_library(hmsbeagle-cpu-sse SHARED
         BeagleCPU4StateSSEImpl.h
@@ -56,6 +58,8 @@ SET_TARGET_PROPERTIES(hmsbeagle-cpu-sse
     SOVERSION "${BEAGLE_PLUGIN_VERSION_EXTENDED}"
 	SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
+
+target_compile_options(hmsbeagle-cpu-sse PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)
 endif(BUILD_SSE)
 
 if(BUILD_SPECTRAL)
@@ -93,4 +97,6 @@ SET_TARGET_PROPERTIES(hmsbeagle-cpu-spectral
     SOVERSION "${BEAGLE_PLUGIN_VERSION_EXTENDED}"
 	SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
+
+target_compile_options(hmsbeagle-cpu-spectral PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)
 endif(BUILD_SPECTRAL)

libhmsbeagle/CPU/EigenDecompositionSpectral.hpp

@@ -87,6 +87,13 @@ void EigenDecompositionSpectral<BEAGLE_CPU_EIGEN_GENERIC>::setEigenDecomposition
     memcpy(tIvec.data(), evec.data(), kMatrixStride * kStateCount * sizeof(REALTYPE));
     transposeInPlace(tIvec.data());
 
+    if (getenv("BEAGLE_DEBUG_EIGEN")) {
+        fprintf(stderr, "[CPU setEigenDecomposition] eigenIndex=%d kStateCount=%d kMatrixStride=%d\n", eigenIndex, kStateCount, kMatrixStride);
+        fprintf(stderr, "[CPU] eval: "); for (int i=0;i<kEigenValuesSize;i++) fprintf(stderr, "%.6f ", (double)eval[i]); fprintf(stderr, "\n");
+        fprintf(stderr, "[CPU] evec:\n"); for (int i=0;i<kStateCount;i++){ for(int j=0;j<kStateCount;j++) fprintf(stderr, "%.6f ", (double)evec[i*kMatrixStride+j]); fprintf(stderr, "\n"); }
+        fprintf(stderr, "[CPU] ivec:\n"); for (int i=0;i<kStateCount;i++){ for(int j=0;j<kStateCount;j++) fprintf(stderr, "%.6f ", (double)ivec[i*kMatrixStride+j]); fprintf(stderr, "\n"); }
+        fflush(stderr);
+    }
     adjointMethodsStorage[eigenIndex] = std::make_shared<AdjointIntegralPlan<REALTYPE>>(
         eigenValuesStorage[eigenIndex].data(), kStateCount, !isComplex);
 }

libhmsbeagle/GPU/BeagleGPUImpl.hpp

@@ -462,6 +462,14 @@ int BeagleGPUImpl<BEAGLE_GPU_GENERIC>::createInstance(int tipCount,
         kFlags |= BEAGLE_FLAG_EIGEN_REAL;
     }
 
+    // kFlags is rebuilt from scratch above (kFlags = 0) and only specific bits are copied
+    // over from preferenceFlags/requirementFlags; BEAGLE_FLAG_SPECTRAL_REPRESENTATION must be
+    // copied too, since BeagleGPUSpectralImpl and the checks just below key off this bit in
+    // kFlags (not off the caller-supplied flags directly).
+    if (preferenceFlags & BEAGLE_FLAG_SPECTRAL_REPRESENTATION || requirementFlags & BEAGLE_FLAG_SPECTRAL_REPRESENTATION) {
+        kFlags |= BEAGLE_FLAG_SPECTRAL_REPRESENTATION;
+    }
+
     if (requirementFlags & BEAGLE_FLAG_INVEVEC_TRANSPOSED || preferenceFlags & BEAGLE_FLAG_INVEVEC_TRANSPOSED)
         kFlags |= BEAGLE_FLAG_INVEVEC_TRANSPOSED;
     else
@@ -497,7 +505,12 @@ int BeagleGPUImpl<BEAGLE_GPU_GENERIC>::createInstance(int tipCount,
     kPartialsSize = kPaddedPatternCount * kPaddedStateCount * kCategoryCount;
     kMatrixSize = kPaddedStateCount * kPaddedStateCount;
 
-    if (kFlags & BEAGLE_FLAG_EIGEN_COMPLEX)
+    // Spectral kernels (kernelsSpectralIfDef*.cu / kernelsSpectral*.cu) always read a
+    // real+imaginary pair per eigenstate (SPECTRAL_EIGENVALS_GPU unconditionally indexes
+    // eigenValues[PADDED_STATE_COUNT + state]) regardless of BEAGLE_FLAG_EIGEN_COMPLEX, so the
+    // device buffer must be sized/copied as complex whenever spectral representation is in use,
+    // otherwise that read runs past the data actually written by setEigenDecomposition.
+    if ((kFlags & BEAGLE_FLAG_EIGEN_COMPLEX) || (kFlags & BEAGLE_FLAG_SPECTRAL_REPRESENTATION))
         kEigenValuesSize = 2 * kPaddedStateCount;
     else
         kEigenValuesSize = kPaddedStateCount;
@@ -4743,6 +4756,7 @@ void BeagleGPUImpl<BEAGLE_GPU_GENERIC>::dispatchPruneSS(GPUPtr s1, GPUPtr s2, GP
                                                          GPUPtr scalingFactors, GPUPtr cumulativeScaling,
                                                          unsigned int startPattern, unsigned int endPattern,
                                                          int rescale, int streamIndex, int waitIndex) {
+    fprintf(stderr, "[DISPATCH] BASE BeagleGPUImpl::dispatchPruneSS called\n"); fflush(stderr);
     kernels->StatesStatesPruningDynamicScaling(s1, s2, p3,
         dMatrices[c1MatIdx], dMatrices[c2MatIdx],
         scalingFactors, cumulativeScaling,

libhmsbeagle/GPU/BeagleGPUSpectralImpl.h

@@ -38,6 +38,12 @@ class BeagleGPUSpectralImpl : public BeagleGPUImpl<Real> {
     GPUPtr  dSpectralDistancesOrigin;
     GPUPtr* dSpectralDistances;
     int*    hEigenIndexForMatrix;
+    /* Per-matrix element stride within dSpectralDistancesOrigin, i.e.
+     * AlignMemOffset(kCategoryCount * sizeof(Real)) / sizeof(Real). Device
+     * alignment padding can make this larger than kCategoryCount, so any code
+     * that indexes into the pooled origin buffer by matrix index (e.g. the
+     * adjoint offset-queue) must multiply by this, not by kCategoryCount. */
+    unsigned int kSpectralDistanceStrideElements;
 
     /* Backward eigenvector buffers for the parent branch in pre-order.
      * dEvecT[ei]  = U   stored row-major (dEvecT[j*S+k] = U[j,k]).

libhmsbeagle/GPU/BeagleGPUSpectralImpl.hpp

@@ -88,15 +88,16 @@ int BeagleGPUSpectralImpl<BEAGLE_GPU_GENERIC>::createInstance(
     GPUInterface* gpuIf = this->gpu;
 
     dSpectralDistances = (GPUPtr*) malloc(sizeof(GPUPtr) * this->kMatrixCount);
-    size_t distStride = this->kCategoryCount * sizeof(Real);
+    size_t distStride = gpuIf->AlignMemOffset(this->kCategoryCount * sizeof(Real));
+    kSpectralDistanceStrideElements = (unsigned int)(distStride / sizeof(Real));
     dSpectralDistancesOrigin = gpuIf->AllocateMemory(this->kMatrixCount * distStride);
     for (int i = 0; i < this->kMatrixCount; i++) {
         dSpectralDistances[i] = gpuIf->CreateSubPointer(dSpectralDistancesOrigin, distStride * i, distStride);
     }
 
     // Backward eigenvector arrays: one S*S block per eigen decomposition.
     int S = this->kPaddedStateCount;
-    size_t matStride = (size_t)S * S * sizeof(Real);
+    size_t matStride = gpuIf->AlignMemOffset((size_t)S * S * sizeof(Real));
     dEvecT  = (GPUPtr*) malloc(sizeof(GPUPtr) * eigenDecompositionCount);
     dIevcT  = (GPUPtr*) malloc(sizeof(GPUPtr) * eigenDecompositionCount);
     dEvecTOrigin  = gpuIf->AllocateMemory(eigenDecompositionCount * matStride);
@@ -191,6 +192,7 @@ void BeagleGPUSpectralImpl<BEAGLE_GPU_GENERIC>::dispatchPruneSS(
         GPUPtr scalingFactors, GPUPtr cumulativeScaling,
         unsigned int startPattern, unsigned int endPattern,
         int rescale, int streamIndex, int waitIndex) {
+    fprintf(stderr, "[DISPATCH] SPECTRAL BeagleGPUSpectralImpl::dispatchPruneSS called\n"); fflush(stderr);
     int ei1 = hEigenIndexForMatrix[c1MatIdx];
     int ei2 = hEigenIndexForMatrix[c2MatIdx];
     this->kernels->StatesStatesPruningSpectral(s1, s2, p3,
@@ -215,7 +217,12 @@ int BeagleGPUSpectralImpl<BEAGLE_GPU_GENERIC>::setEigenDecomposition(
     const int SC = this->kStateCount;
     const int SS = S * S;
     // kEigenValuesSize is private to BeagleGPUImpl; recompute the same way it does.
-    const int eigenValuesSize = (this->kFlags & BEAGLE_FLAG_EIGEN_COMPLEX) ? 2 * S : S;
+    // Spectral kernels always read a real+imaginary pair per eigenstate regardless of
+    // BEAGLE_FLAG_EIGEN_COMPLEX (see matching comment/fix in BeagleGPUImpl::createInstance),
+    // so this local copy must widen for BEAGLE_FLAG_SPECTRAL_REPRESENTATION too, or the
+    // imaginary half of the device buffer is left uninitialized.
+    const int eigenValuesSize = ((this->kFlags & BEAGLE_FLAG_EIGEN_COMPLEX) ||
+                                  (this->kFlags & BEAGLE_FLAG_SPECTRAL_REPRESENTATION)) ? 2 * S : S;
 
     // Forward:  dEvec[row*S+col]  = U[col,row];    dIevc[row*S+col]  = U^-1[col,row]
     // Backward: dEvecT[row*S+col] = U[row,col];    dIevcT[row*S+col] = U^-1[row,col]
@@ -280,6 +287,13 @@ int BeagleGPUSpectralImpl<BEAGLE_GPU_GENERIC>::setEigenDecomposition(
             Eval[S + i] = (Real) inEigenValues[SC + i];
     }
 
+    if (getenv("BEAGLE_DEBUG_EIGEN")) {
+        fprintf(stderr, "[GPU setEigenDecomposition] eigenIndex=%d S=%d SC=%d\n", eigenIndex, S, SC);
+        fprintf(stderr, "[GPU] Eval: "); for (int i=0;i<eigenValuesSize;i++) fprintf(stderr, "%.6f ", (double)Eval[i]); fprintf(stderr, "\n");
+        fprintf(stderr, "[GPU] Evec:\n"); for (int i=0;i<S;i++){ for(int j=0;j<S;j++) fprintf(stderr, "%.6f ", (double)Evec[i*S+j]); fprintf(stderr, "\n"); }
+        fprintf(stderr, "[GPU] Ievc:\n"); for (int i=0;i<S;i++){ for(int j=0;j<S;j++) fprintf(stderr, "%.6f ", (double)Ievc[i*S+j]); fprintf(stderr, "\n"); }
+        fflush(stderr);
+    }
     GPUInterface* gpuIf = this->gpu;
     gpuIf->MemcpyHostToDevice(this->dEvec[eigenIndex],        Evec,  sizeof(Real) * SS);
     gpuIf->MemcpyHostToDevice(this->dIevc[eigenIndex],        Ievc,  sizeof(Real) * SS);
@@ -599,7 +613,7 @@ int BeagleGPUSpectralImpl<BEAGLE_GPU_GENERIC>::calculateAdjointCrossProducts(
             rec[3] = (unsigned int)ei * SS;
             rec[4] = (unsigned int)ei * this->getEvecStrideElements();
             rec[5] = (unsigned int)ei * this->getEvalStrideElements();
-            rec[6] = (unsigned int)matIdx * this->kCategoryCount;
+            rec[6] = (unsigned int)matIdx * kSpectralDistanceStrideElements;
             rec[7] = (unsigned int)ei * SS;
             rec[8] = isAllReal ? 1u : 0u;
         }
@@ -625,13 +639,24 @@ int BeagleGPUSpectralImpl<BEAGLE_GPU_GENERIC>::calculateAdjointCrossProducts(
     }
 
     /* Download gradient to host.  The first eigen decomp's gradient is the
-     * primary output.  Convert Real→double. */
+     * primary output.  Convert Real→double.
+     *
+     * The device buffer is laid out S×S with S=kPaddedStateCount (padding
+     * rows/cols beyond kStateCount hold unused/garbage values from the
+     * padding eigenstates), but callers (matching BeagleCPUImpl's
+     * convention) allocate outSumDerivatives sized kStateCount*kStateCount.
+     * Copying the full padded SS run would overflow that buffer whenever
+     * kPaddedStateCount > kStateCount — only copy the top-left
+     * kStateCount×kStateCount submatrix, row by row with the correct
+     * strides on each side. */
     if (count > 0 && outSumDerivatives != NULL) {
         const int ei0 = hEigenIndexForMatrix[eigenIndices[0]];
+        const int SC = this->kStateCount;
         Real* hGrad = (Real*) gpuIf->CallocHost(sizeof(Real), SS);
         gpuIf->MemcpyDeviceToHost(hGrad, dGradient[ei0], SS * sizeof(Real));
-        for (int k = 0; k < SS; k++)
-            outSumDerivatives[k] = (double)hGrad[k];
+        for (int i = 0; i < SC; i++)
+            for (int j = 0; j < SC; j++)
+                outSumDerivatives[i * SC + j] = (double)hGrad[i * S + j];
         gpuIf->FreeHostMemory(hGrad);
     }
 

libhmsbeagle/GPU/CMake_CUDA/CMakeLists.txt

@@ -62,3 +62,5 @@ SET_TARGET_PROPERTIES(hmsbeagle-cuda
     SOVERSION "${BEAGLE_PLUGIN_VERSION_EXTENDED}"
 	SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
+
+target_compile_options(hmsbeagle-cuda PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)

libhmsbeagle/GPU/CMake_CUDASpectral/CMakeLists.txt

@@ -70,3 +70,5 @@ SET_TARGET_PROPERTIES(hmsbeagle-cuda-spectral
     SOVERSION "${BEAGLE_PLUGIN_VERSION_EXTENDED}"
     SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
+
+target_compile_options(hmsbeagle-cuda-spectral PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)

libhmsbeagle/GPU/CMake_OpenCL/CMakeLists.txt

@@ -53,3 +53,5 @@ SET_TARGET_PROPERTIES(hmsbeagle-opencl
     SOVERSION "${BEAGLE_PLUGIN_VERSION_EXTENDED}"
 	SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
+
+target_compile_options(hmsbeagle-opencl PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)

libhmsbeagle/GPU/CMake_OpenCLSpectral/CMakeLists.txt

@@ -62,3 +62,5 @@ SET_TARGET_PROPERTIES(hmsbeagle-opencl-spectral
     SOVERSION "${BEAGLE_PLUGIN_VERSION_EXTENDED}"
     SUFFIX "${BEAGLE_PLUGIN_SUFFIX}"
     )
+
+target_compile_options(hmsbeagle-opencl-spectral PRIVATE -fvisibility=hidden -fvisibility-inlines-hidden)

libhmsbeagle/GPU/CUDAPlugin.cpp

@@ -77,7 +77,7 @@ CUDAPlugin::~CUDAPlugin()
 
 
 extern "C" {
-void* plugin_init(void){
+__attribute__((visibility("default"))) void* plugin_init(void){
 	return new beagle::gpu::CUDAPlugin();
 }
 }

libhmsbeagle/GPU/CUDASpectralPlugin.cpp

@@ -74,7 +74,7 @@ CUDASpectralPlugin::~CUDASpectralPlugin()
 
 
 extern "C" {
-void* plugin_init(void) {
+__attribute__((visibility("default"))) void* plugin_init(void) {
     return new beagle::gpu::CUDASpectralPlugin();
 }
 }