From 03f062d9235c4e66373ceb989b80b14391264816 Mon Sep 17 00:00:00 2001 From: fnachon Date: Thu, 16 Jul 2026 18:05:45 +0200 Subject: [PATCH 1/3] Add Metal/MPS GPU backend for the ungapped prefilter on Apple Silicon Implements a Metal-based counterpart to lib/libmarv (CUDA/HIP) behind a new ENABLE_MPS CMake option, so --gpu 1 works on Apple Silicon Macs. Scope is the default ungapped prefilter (GAPLESS) only; the gapped Smith-Waterman GPU path remains CUDA/HIP-only. - lib/libmarv-mps/: Metal implementation of the Marv interface (marv.h), reusing it unchanged so GpuUtil/gpuserver/ungappedprefilter need no backend-specific logic. Kernel design: one simdgroup cooperates per alignment (registers-in-threadgroup-memory DP state, coalesced PSSM layout), with work-stealing dispatch across a shared atomic counter so skewed sequence-length databases don't stall on the longest outliers. submitScan()/collectScan() split lets the query loop overlap GPU execution with building the next query's profile. - marv.h: additive submitScan()/collectScan() declarations for the async split; CUDA/HIP is unaffected since nothing calls them under that build. - CMakeLists.txt, src/CMakeLists.txt: ENABLE_MPS option, mutually exclusive with ENABLE_CUDA. - GpuUtil.cpp, Parameters.cpp, gpuserver.cpp, ungappedprefilter.cpp: extend existing HAVE_CUDA guards to HAVE_MPS; CUDA/HIP code paths are unchanged. - README.md: build/usage note for the new backend. Validated on an Apple M5 Max: byte-identical results against a scalar CPU reference and against MMseqs2's own CPU prefilter (modulo the CPU path's intentional 8-bit score saturation), across synthetic boundary-length cases and real searches up to 153.7M sequences (UniRef90). GPU search against UniRef90 measured ~8.5x faster than the CPU path's k-mer prefilter for the same query on this database size. Co-Authored-By: Claude Sonnet 5 --- CMakeLists.txt | 11 + README.md | 12 + lib/libmarv-mps/CMakeLists.txt | 17 + lib/libmarv-mps/marv_mps.mm | 534 +++++++++++++++++++++++++ lib/libmarv/src/marv.h | 10 + src/CMakeLists.txt | 5 + src/commons/GpuUtil.cpp | 2 +- src/commons/Parameters.cpp | 4 +- src/prefiltering/ungappedprefilter.cpp | 191 +++++---- src/util/gpuserver.cpp | 4 +- 10 files changed, 714 insertions(+), 76 deletions(-) create mode 100644 lib/libmarv-mps/CMakeLists.txt create mode 100644 lib/libmarv-mps/marv_mps.mm diff --git a/CMakeLists.txt b/CMakeLists.txt index 04dc1dfdc..0a0c35c02 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,6 +21,7 @@ set(USE_SYSTEM_ZSTD 0 CACHE BOOL "Use zstd provided by system instead of bundled set(IGNORE_RUST_VERSION 0 CACHE BOOL "Ignore Rust version check") set(ENABLE_HIP 0 CACHE BOOL "Enable HIP") set(ENABLE_CUDA 0 CACHE BOOL "Enable CUDA") +set(ENABLE_MPS 0 CACHE BOOL "Enable Metal/MPS GPU backend (Apple Silicon)") set(FORCE_STATIC_DEPS 0 CACHE BOOL "Force static linking of deps") set(MMSEQS_INT64_IDS 0 CACHE BOOL "Use 64-bit internal sequence/database IDs") @@ -28,6 +29,10 @@ if(ENABLE_HIP) set(ENABLE_CUDA 1) endif() +if(ENABLE_CUDA AND ENABLE_MPS) + message(FATAL_ERROR "ENABLE_CUDA and ENABLE_MPS are mutually exclusive") +endif() + if(FORCE_STATIC_DEPS) if(ENABLE_CUDA) set(CMAKE_FIND_LIBRARY_SUFFIXES .a .so CACHE INTERNAL "" FORCE) @@ -310,6 +315,12 @@ if (ENABLE_CUDA) include_directories(lib/libmarv/src) add_subdirectory(lib/libmarv/src EXCLUDE_FROM_ALL) set_target_properties(marv PROPERTIES POSITION_INDEPENDENT_CODE ON) +elseif (ENABLE_MPS) + # Metal/MPS backend for Apple Silicon; reuses the same marv.h interface + # as lib/libmarv so downstream code (GpuUtil, gpuserver, ungappedprefilter) + # needs no changes beyond the HAVE_MPS compile definition. + include_directories(lib/libmarv/src) + add_subdirectory(lib/libmarv-mps EXCLUDE_FROM_ALL) endif () add_subdirectory(lib/alp) diff --git a/README.md b/README.md index 87f09d98b..3753039a7 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,18 @@ To perform searches using GPU acceleration, you can add the `--gpu` flag to `eas mmseqs createdb examples/DB.fasta targetDB mmseqs makepaddedseqdb targetDB targetDB_padded mmseqs easy-search examples/QUERY.fasta targetDB_padded alnRes.m8 tmp --gpu 1 + +#### Apple Silicon (Metal/MPS) GPU support + +> [!NOTE] +> This is a local, unofficial backend (not part of the upstream project or the precompiled binaries) that ports GPU-accelerated search to Apple Silicon via Metal, as an alternative to the CUDA backend above. It currently only accelerates the default ungapped prefilter (`--prefilter-mode 0`, the same mode used by `--gpu 1` above); the gapped Smith-Waterman GPU kernel (`--prefilter-mode 2`) is CUDA/HIP-only for now. + +Build with `-DENABLE_MPS=1` instead of `-DENABLE_CUDA=1` (the two are mutually exclusive): + + cmake -DENABLE_MPS=1 -DHAVE_ARM8=1 -DCMAKE_BUILD_TYPE=Release .. + make -j + +Usage is identical to the CUDA workflow above (`makepaddedseqdb` + `--gpu 1`); no macOS-specific flags are needed. See [lib/libmarv-mps/marv_mps.mm](lib/libmarv-mps/marv_mps.mm) for the implementation and design notes. The `databases` workflow provides download and setup procedures for many public reference databases, such as the Uniref, NR, NT, PFAM and many more (see [Downloading databases](https://github.com/soedinglab/mmseqs2/wiki#downloading-databases)). For example, to download and search against a database containing the Swiss-Prot reference proteins run: diff --git a/lib/libmarv-mps/CMakeLists.txt b/lib/libmarv-mps/CMakeLists.txt new file mode 100644 index 000000000..0f6886ccd --- /dev/null +++ b/lib/libmarv-mps/CMakeLists.txt @@ -0,0 +1,17 @@ +cmake_minimum_required(VERSION 3.18) + +if (NOT APPLE) + message(FATAL_ERROR "libmarv-mps (Metal/MPS GPU backend) requires macOS") +endif () + +enable_language(OBJCXX) + +add_library(marv STATIC marv_mps.mm) +set_source_files_properties(marv_mps.mm PROPERTIES LANGUAGE OBJCXX) +target_compile_features(marv PUBLIC cxx_std_17) +target_compile_options(marv PRIVATE -fobjc-arc) +set_target_properties(marv PROPERTIES POSITION_INDEPENDENT_CODE ON) + +find_library(METAL_FRAMEWORK Metal REQUIRED) +find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED) +target_link_libraries(marv PUBLIC ${METAL_FRAMEWORK} ${FOUNDATION_FRAMEWORK}) diff --git a/lib/libmarv-mps/marv_mps.mm b/lib/libmarv-mps/marv_mps.mm new file mode 100644 index 000000000..ab10e00af --- /dev/null +++ b/lib/libmarv-mps/marv_mps.mm @@ -0,0 +1,534 @@ +// Metal/MPS implementation of the Marv GPU-prefilter interface declared in +// lib/libmarv/src/marv.h. This is the Apple Silicon counterpart to +// lib/libmarv (CUDA/HIP): it implements the same "gapless" (ungapped) +// local-alignment scan used to accelerate MMseqs2's prefilter stage, so it +// slots into the existing GpuUtil.cpp / gpuserver.cpp / ungappedprefilter.cpp +// pipeline without changes to their control flow. +// +// Scope: only AlignmentType::GAPLESS is implemented (the default prefilter +// path). The gapped Smith-Waterman kernel in libmarv (used by +// --prefilter-mode 2) has no Metal equivalent yet. +// +// Algorithm: H[i][j] = max(0, H[i-1][j-1] + score(query[i], subject[j])), +// boundary zero -- the same no-gap recurrence as libmarv's GAPLESS kernel +// (pssmkernels_gapless.cuh). Subject bytes are folded with `& 31` to +// collapse MMseqs2's soft-masked (lowercase, +32) residue encoding back to +// the base alphabet, mirroring libmarv's CaseSensitive_to_CaseInsensitive +// conversion (see convert.cuh). +// +// Parallelism (v2, simdgroup-cooperative): a first version of this kernel +// used one GPU thread per subject sequence, with its O(queryLen) DP state in +// `device` (global) memory. That was measured to be ~12x slower than +// MMseqs2's CPU prefilter on a real, length-skewed database (lengths 9 to +// 8083 residues, median 347): a handful of very long sequences dominated +// each dispatch's wall-clock while thousands of short-sequence threads sat +// idle, and every DP step round-tripped through slow global memory. +// +// This version instead mirrors libmarv's actual design: one Metal simdgroup +// (32 lanes, lockstep -- Apple's equivalent of a CUDA warp) cooperates on a +// single alignment. Each lane owns a contiguous slice of the query in fast +// `threadgroup` memory; per subject residue, one simd_shuffle_up carries the +// single boundary value from lane L's slice to lane L+1's (the only +// cross-lane dependency the recurrence has -- see GaplessPSSMState::relax / +// shuffleScores in pssmkernels_gapless.cuh for the CUDA original), then all +// 32 lanes update their slice in parallel. Unlike libmarv, the slice length +// is a runtime loop over dynamically-sized threadgroup memory rather than a +// fixed-size register array selected from a family of compiled kernel +// variants -- this trades some peak throughput for supporting arbitrary +// query lengths with a single kernel. A further simdgroup dispatched per +// database entry would still leave the same long-tail scheduling problem, +// so simdgroups instead pull work from a shared atomic counter until the +// database is exhausted: idle simdgroups immediately pick up the next +// (typically short) sequence instead of waiting on whichever simdgroup drew +// the longest one. +// +// PSSM layout (v3, coalesced): with lanes owning *contiguous* query slices, +// reading pssm[s][chunkStart+li] at loop step `li` means lane L reads +// address L*chunkSize+li -- lanes are chunkSize apart, not adjacent, so the +// 32 lanes' simultaneous reads hit scattered cache lines instead of one +// coalesced burst. submitScan() uploads each residue row pre-transposed +// into chunkSize*32 "column-major" blocks (row[li*32+lane] instead of +// row[lane*chunkSize+li]), matching libmarv's own smem layout trick (see +// SmemIndexCalculator in pssmkernels_gapless.cuh) -- the 32 lanes now read +// 32 consecutive bytes each step. The DP state in threadgroup memory (Hall) +// stays untransposed since it's private per simdgroup, not shared. Measured +// impact was modest (0-20% depending on query length, <1% in aggregate on a +// real query mix) -- Apple Silicon's unified GPU memory system apparently +// already absorbs the scattered pattern reasonably well via its cache. +// +// Register tiling -- tried, reverted: a v4 attempt moved the DP state out +// of threadgroup memory into a fixed-size per-lane register array (32 +// int32 registers/lane, covering queries up to 1024 residues via 4 tiles of +// 256, falling back to the threadgroup-memory kernel above for longer +// queries), the direct Metal analogue of libmarv's register-resident, +// border_in/border_out tile-chained kernels. It was verified correct -- +// byte-identical output against the threadgroup-memory kernel across 25 +// boundary-length test cases (1, 2, 31-33, 63-65, ..., 1023-1025, 2000) and +// the full 500-query/20k-sequence real benchmark -- but measured ~27% +// *slower* on the queries it targeted (8.36s vs an estimated 6.59s for the +// same 457 queries under the threadgroup-memory kernel). +// +// Checked with Instruments (Metal System Trace, via `xcrun xctrace record +// --template "Metal System Trace"`) rather than guessing: the +// graphics-compiler-spill-events table showed zero spill events for either +// kernel, ruling out register spilling. metal-shader-profiler-shader-list +// showed the actual difference -- gaplessAlignRegisterTiled compiled to +// 2472 bytes of GPU code vs. gaplessAlignCooperative's 1006, 2.46x more, +// almost certainly from the compiler fully unrolling the compile-time-sized +// tile*register loops (4 tiles x 8 registers) plus the per-register +// `globalPos < qlen` bounds check and the extra cross-tile border shuffle +// bookkeeping. +// +// Follow-up: is it the unrolling (code size) or the bounds-check/shuffle +// overhead that dominates? Tested by making the inner register loop bound a +// runtime Params field instead of the compile-time constant kRegs, so the +// compiler could no longer statically unroll it. Code size did drop, to +// 1484 bytes (vs. 2472 unrolled, 1006 cooperative) -- but wall-clock got +// *worse*, not better: 15.3s vs. 11.7s (unrolled) vs. 9.9s (cooperative) +// on the same 500-query benchmark. graphics-compiler-spill-events explained +// why: this version spilled 144 bytes/thread on essentially every dispatch, +// where the unrolled version spilled nothing. GPU registers can't be +// dynamically indexed -- a loop bound the compiler can't resolve at compile +// time forces it to either unroll anyway (code bloat) or address the array +// through memory (spilling), and spilling turned out to cost more than the +// code bloat it replaced. So it's not really "unrolling vs. bounds-check +// overhead" -- both register-tiled variants lose to the threadgroup-memory +// kernel, for complementary reasons, because there's no way to keep a +// variably-sized per-lane array genuinely register-resident on this +// hardware. The threadgroup-memory kernel wins by making the memory access +// explicit and coalesced instead of fighting the register model. This +// looks like the ceiling for this kernel design; a further improvement +// would need a fundamentally different approach (e.g. libmarv's actual +// fixed per-(group_size,numRegs)-config kernel family with host-side +// config selection per query length), not another variant of this one. + +#import +#import + +#include "marv.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +// Upper bound on how many simdgroups cooperate within one threadgroup. +// Actual value is chosen at dispatch time from the device's threadgroup +// memory budget; this just avoids pointlessly large threadgroups for short +// queries. +constexpr uint32_t kMaxSimdgroupsPerThreadgroup = 8; +// Upper bound on total simdgroups launched across the whole grid; with +// work-stealing this only needs to be "enough to saturate the GPU", not +// one per database entry. +constexpr uint32_t kMaxTotalSimdgroups = 4096; + +const char* kKernelSource = R"METAL( +#include +using namespace metal; + +struct Params { + uint32_t dbEntries; + uint32_t queryLen; +}; + +kernel void gaplessAlignCooperative( + constant Params& params [[buffer(0)]], + device const int8_t* pssm [[buffer(1)]], + device const uint8_t* dbData [[buffer(2)]], + device const uint64_t* offsets [[buffer(3)]], + device const uint32_t* lengths [[buffer(4)]], + device atomic_uint* workCounter [[buffer(5)]], + device int32_t* results [[buffer(6)]], + threadgroup int32_t* Hall [[threadgroup(0)]], + uint simdLane [[thread_index_in_simdgroup]], + uint simdGroupId [[simdgroup_index_in_threadgroup]]) +{ + const uint32_t qlen = params.queryLen; + threadgroup int32_t* myH = Hall + (size_t)simdGroupId * qlen; + + // Each of the 32 lanes owns a contiguous slice of the query. `pssm` rows + // are pre-transposed into chunkSize*32 blocks so that step `li` reads + // row[li*32 + lane] -- 32 consecutive bytes -- instead of the scattered + // row[lane*chunkSize + li] a plain row-major layout would give. + const uint32_t chunkSize = (qlen + 31u) / 32u; + const uint32_t chunkStart = min(simdLane * chunkSize, qlen); + const uint32_t chunkEnd = min(chunkStart + chunkSize, qlen); + const uint32_t paddedQueryLen = chunkSize * 32u; + + while (true) { + // Lane 0 claims the next database entry for this simdgroup; the + // claimed index is broadcast so every lane agrees on when to stop + // (required for the simd_shuffle_up calls below to stay uniform). + uint32_t work = 0; + if (simdLane == 0) { + work = atomic_fetch_add_explicit(workCounter, 1u, memory_order_relaxed); + } + work = simd_broadcast_first(work); + if (work >= params.dbEntries) { + break; + } + + for (uint32_t i = chunkStart; i < chunkEnd; i++) { + myH[i] = 0; + } + + const uint64_t start = offsets[work]; + const uint32_t len = lengths[work]; + + int32_t best = 0; + for (uint32_t j = 0; j < len; j++) { + const uint8_t s = dbData[start + j] & 31; + device const int8_t* row = pssm + (size_t)s * paddedQueryLen; + + // The only cross-lane dependency: H[i-1][j-1] for this lane's + // first position is H[chunkEnd-1][j-1] from the lane below, + // i.e. that lane's last slot *before* it gets overwritten this + // step. Read-then-shuffle happens before any lane writes to + // `myH` for step j, so this is race-free without a barrier. + const int32_t myLastOld = (chunkEnd > chunkStart) ? myH[chunkEnd - 1] : 0; + int32_t borderIn = simd_shuffle_up(myLastOld, 1); + if (simdLane == 0) { + borderIn = 0; // query position -1 boundary + } + + int32_t carry = borderIn; + const uint32_t numLocal = chunkEnd - chunkStart; + for (uint32_t li = 0; li < numLocal; li++) { + const uint32_t i = chunkStart + li; + const int32_t oldVal = myH[i]; + const int32_t val = max(carry + int32_t(row[li * 32u + simdLane]), 0); + myH[i] = val; + best = max(best, val); + carry = oldVal; + } + } + + const int32_t groupBest = simd_max(best); + if (simdLane == 0) { + results[work] = groupBest; + } + } +} +)METAL"; + +// Score/index pair used by the bounded top-K selection in collectScan(). +struct ScoreIdx { + int32_t score; + uint32_t id; +}; + +} // namespace + +struct MarvDb { + id data = nil; + id offsets = nil; + id lengths = nil; + size_t dbEntries = 0; +}; + +struct MarvContext { + id device = nil; + id queue = nil; + id pipeline = nil; + size_t maxThreadgroupMemory = 0; + + id resultsBuf = nil; // dbEntries int32 scores, reused across scans + id workCounterBuf = nil; // single atomic_uint, reset before each scan + + id pssmBuf = nil; // reused, grown on demand instead of realloc'd every scan + size_t pssmCapacityBytes = 0; + + // In-flight GPU work started by submitScan(), consumed by collectScan(). + id pendingCmdBuf = nil; + size_t pendingTotalCells = 0; + std::chrono::high_resolution_clock::time_point pendingSubmitTime; + + std::vector topKHeap; // reused scratch for collectScan()'s top-K selection + + MarvDb* activeDb = nullptr; + size_t maxResults = 0; +}; + +Marv::Marv(size_t dbEntries, int alphabetSize, int maxSeqLength, size_t maxSeqs, Marv::AlignmentType alignmentType) + : dbEntries(dbEntries), alphabetSize(alphabetSize), dbmanager(NULL), alignmentType(alignmentType) { + (void)maxSeqLength; + + if (alignmentType != AlignmentType::GAPLESS) { + throw std::runtime_error( + "Marv (Metal/MPS backend): only AlignmentType::GAPLESS is currently implemented. " + "Gapped Smith-Waterman GPU alignment is not yet ported to Apple Silicon -- " + "use the default --prefilter-mode 0 with --gpu 1."); + } + + MarvContext* ctx = new MarvContext(); + ctx->device = MTLCreateSystemDefaultDevice(); + if (!ctx->device) { + delete ctx; + throw std::runtime_error("Marv (MPS): no Metal device available"); + } + ctx->maxThreadgroupMemory = [ctx->device maxThreadgroupMemoryLength]; + + MTLCompileOptions* compileOptions = [MTLCompileOptions new]; + compileOptions.languageVersion = MTLLanguageVersion3_0; + + NSError* error = nil; + id library = [ctx->device newLibraryWithSource:[NSString stringWithUTF8String:kKernelSource] + options:compileOptions + error:&error]; + if (!library) { + std::string msg = "Marv (MPS): shader compile failed: "; + msg += error ? [[error localizedDescription] UTF8String] : "unknown error"; + delete ctx; + throw std::runtime_error(msg); + } + id fn = [library newFunctionWithName:@"gaplessAlignCooperative"]; + ctx->pipeline = [ctx->device newComputePipelineStateWithFunction:fn error:&error]; + if (!ctx->pipeline) { + std::string msg = "Marv (MPS): pipeline creation failed: "; + msg += error ? [[error localizedDescription] UTF8String] : "unknown error"; + delete ctx; + throw std::runtime_error(msg); + } + ctx->queue = [ctx->device newCommandQueue]; + + ctx->maxResults = std::min(maxSeqs, dbEntries > 0 ? dbEntries : maxSeqs); + ctx->resultsBuf = [ctx->device newBufferWithLength:std::max((size_t)1, dbEntries) * sizeof(int32_t) + options:MTLResourceStorageModeShared]; + ctx->workCounterBuf = [ctx->device newBufferWithLength:sizeof(uint32_t) + options:MTLResourceStorageModeShared]; + ctx->topKHeap.reserve(ctx->maxResults); + + cudasw = static_cast(ctx); +} + +Marv::~Marv() { + MarvContext* ctx = static_cast(cudasw); + delete ctx; +} + +std::vector Marv::getDeviceIds() { + id device = MTLCreateSystemDefaultDevice(); + if (!device) { + return {}; + } + return {0}; +} + +void* Marv::loadDb(char* data, size_t* offset, int32_t* length, size_t dbByteSize) { + MarvContext* ctx = static_cast(cudasw); + + MarvDb* db = new MarvDb(); + db->dbEntries = dbEntries; + db->data = [ctx->device newBufferWithBytes:data length:dbByteSize options:MTLResourceStorageModeShared]; + if (!db->data) { + delete db; + throw std::runtime_error("Marv (MPS): failed to allocate GPU buffer for database"); + } + + std::vector offsets64(dbEntries); + std::vector lengths32(dbEntries); + for (size_t i = 0; i < dbEntries; i++) { + offsets64[i] = static_cast(offset[i]); + lengths32[i] = static_cast(length[i]); + } + db->offsets = [ctx->device newBufferWithBytes:offsets64.data() + length:offsets64.size() * sizeof(uint64_t) + options:MTLResourceStorageModeShared]; + db->lengths = [ctx->device newBufferWithBytes:lengths32.data() + length:lengths32.size() * sizeof(uint32_t) + options:MTLResourceStorageModeShared]; + return static_cast(db); +} + +void* Marv::loadDb(char* /*data*/, size_t /*dbByteSize*/, void* /*otherdb*/) { + throw std::runtime_error("Marv (MPS): cross-process external DB sharing is not supported by the Metal backend"); +} + +void Marv::setDb(void* dbhandle) { + MarvContext* ctx = static_cast(cudasw); + ctx->activeDb = static_cast(dbhandle); +} + +void Marv::setDbWithAllocation(void* /*dbhandle*/, const std::string& /*allocationinfo*/) { + throw std::runtime_error("Marv (MPS): cross-process external DB sharing is not supported by the Metal backend"); +} + +std::string Marv::getDbMemoryHandle() { + throw std::runtime_error("Marv (MPS): cross-process external DB sharing is not supported by the Metal backend"); +} + +void Marv::printInfo() { + MarvContext* ctx = static_cast(cudasw); + fprintf(stderr, "Marv (Metal/MPS): device=%s dbEntries=%zu alphabetSize=%d\n", + [[ctx->device name] UTF8String], dbEntries, alphabetSize); +} + +void Marv::prefetch() { + // No-op: Apple Silicon uses unified memory, so there is no separate + // device memory to migrate data into ahead of time. +} + +void Marv::startTimer() {} +void Marv::stopTimer() {} + +// sequence must be encoded +void Marv::submitScan(const char* /*sequence*/, size_t sequenceLength, int8_t* pssm) { + MarvContext* ctx = static_cast(cudasw); + if (ctx->activeDb == nullptr) { + throw std::runtime_error("Marv (MPS): submitScan() called before setDb()"); + } + if (ctx->pendingCmdBuf != nil) { + throw std::runtime_error("Marv (MPS): submitScan() called while a previous scan is still pending -- call collectScan() first"); + } + MarvDb* db = ctx->activeDb; + const uint32_t queryLen = static_cast(sequenceLength); + + const size_t bytesPerSimdgroup = (size_t)queryLen * sizeof(int32_t); + if (bytesPerSimdgroup > ctx->maxThreadgroupMemory) { + throw std::runtime_error( + "Marv (MPS): query too long for the Metal backend (" + std::to_string(queryLen) + + " residues needs " + std::to_string(bytesPerSimdgroup) + " bytes of threadgroup memory, device limit is " + + std::to_string(ctx->maxThreadgroupMemory) + ")"); + } + uint32_t simdgroupsPerTG = static_cast( + std::min(kMaxSimdgroupsPerThreadgroup, ctx->maxThreadgroupMemory / bytesPerSimdgroup)); + simdgroupsPerTG = std::max(1u, simdgroupsPerTG); + + // Reuse the pssm buffer across calls instead of allocating a fresh + // MTLBuffer every query -- it's a small copy, but resource creation has + // real fixed overhead that adds up over hundreds/thousands of queries. + // + // Uploaded pre-transposed to match the kernel's per-lane striping (see + // the file header comment): row `r`'s chunkSize*32 block holds + // dst[li*32+lane] = pssm[r][lane*chunkSize+li], so that the 32 lanes' + // simultaneous reads at loop step `li` land on 32 consecutive bytes + // instead of being chunkSize apart. Padding past each lane's actual + // chunkEnd is never read by the kernel, so it's left uninitialized. + const uint32_t chunkSize = (queryLen + 31u) / 32u; + const uint32_t paddedQueryLen = chunkSize * 32u; + const size_t pssmBytesNeeded = (size_t)alphabetSize * paddedQueryLen * sizeof(int8_t); + if (pssmBytesNeeded > ctx->pssmCapacityBytes) { + ctx->pssmBuf = [ctx->device newBufferWithLength:pssmBytesNeeded options:MTLResourceStorageModeShared]; + if (!ctx->pssmBuf) { + throw std::runtime_error("Marv (MPS): failed to allocate pssm buffer"); + } + ctx->pssmCapacityBytes = pssmBytesNeeded; + } + int8_t* dst = static_cast([ctx->pssmBuf contents]); + for (uint32_t r = 0; r < (uint32_t)alphabetSize; r++) { + const int8_t* srcRow = pssm + (size_t)r * queryLen; + int8_t* dstRow = dst + (size_t)r * paddedQueryLen; + for (uint32_t lane = 0; lane < 32u; lane++) { + const uint32_t chunkStart = std::min(lane * chunkSize, queryLen); + const uint32_t chunkEnd = std::min(chunkStart + chunkSize, queryLen); + for (uint32_t li = 0; li < chunkEnd - chunkStart; li++) { + dstRow[li * 32u + lane] = srcRow[chunkStart + li]; + } + } + } + + std::memset([ctx->workCounterBuf contents], 0, sizeof(uint32_t)); + + struct { + uint32_t dbEntries; + uint32_t queryLen; + } params{static_cast(db->dbEntries), queryLen}; + + ctx->pendingTotalCells = (size_t)queryLen * (size_t)[db->data length]; + ctx->pendingSubmitTime = std::chrono::high_resolution_clock::now(); + + if (db->dbEntries == 0) { + ctx->pendingCmdBuf = nil; // nothing to dispatch; collectScan() handles this + return; + } + + const uint32_t totalSimdgroups = static_cast(std::min(db->dbEntries, kMaxTotalSimdgroups)); + const uint32_t numThreadgroups = std::max(1u, (totalSimdgroups + simdgroupsPerTG - 1) / simdgroupsPerTG); + + id cmdBuf = [ctx->queue commandBuffer]; + id encoder = [cmdBuf computeCommandEncoder]; + [encoder setComputePipelineState:ctx->pipeline]; + [encoder setBytes:¶ms length:sizeof(params) atIndex:0]; + [encoder setBuffer:ctx->pssmBuf offset:0 atIndex:1]; + [encoder setBuffer:db->data offset:0 atIndex:2]; + [encoder setBuffer:db->offsets offset:0 atIndex:3]; + [encoder setBuffer:db->lengths offset:0 atIndex:4]; + [encoder setBuffer:ctx->workCounterBuf offset:0 atIndex:5]; + [encoder setBuffer:ctx->resultsBuf offset:0 atIndex:6]; + [encoder setThreadgroupMemoryLength:simdgroupsPerTG * bytesPerSimdgroup atIndex:0]; + + [encoder dispatchThreadgroups:MTLSizeMake(numThreadgroups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(simdgroupsPerTG * 32, 1, 1)]; + [encoder endEncoding]; + [cmdBuf commit]; + + ctx->pendingCmdBuf = cmdBuf; // not waited on -- collectScan() does that +} + +Marv::Stats Marv::collectScan(Result* results) { + MarvContext* ctx = static_cast(cudasw); + if (ctx->activeDb == nullptr) { + throw std::runtime_error("Marv (MPS): collectScan() called before setDb()"); + } + MarvDb* db = ctx->activeDb; + + if (ctx->pendingCmdBuf != nil) { + [ctx->pendingCmdBuf waitUntilCompleted]; + if (ctx->pendingCmdBuf.status == MTLCommandBufferStatusError) { + std::string msg = "Marv (MPS): command buffer failed: "; + msg += ctx->pendingCmdBuf.error ? [[ctx->pendingCmdBuf.error localizedDescription] UTF8String] : "unknown error"; + ctx->pendingCmdBuf = nil; + throw std::runtime_error(msg); + } + ctx->pendingCmdBuf = nil; + } else if (db->dbEntries != 0) { + throw std::runtime_error("Marv (MPS): collectScan() called without a matching submitScan()"); + } + + auto timeEnd = std::chrono::high_resolution_clock::now(); + double seconds = std::chrono::duration(timeEnd - ctx->pendingSubmitTime).count(); + + const int32_t* scores = static_cast([ctx->resultsBuf contents]); + + // Bounded top-K selection: keep a min-heap of at most maxResults + // entries instead of sorting all dbEntries scores every query. Cheap + // even for very large databases since the heap never grows past + // maxResults (typically a few hundred). + const size_t maxResults = std::min(ctx->maxResults, db->dbEntries); + std::vector& heap = ctx->topKHeap; + heap.clear(); + auto heapCmp = [](const ScoreIdx& a, const ScoreIdx& b) { return a.score > b.score; }; // min-heap by score + for (size_t i = 0; i < db->dbEntries; i++) { + const int32_t sc = scores[i]; + if (heap.size() < maxResults) { + heap.push_back(ScoreIdx{sc, static_cast(i)}); + std::push_heap(heap.begin(), heap.end(), heapCmp); + } else if (sc > heap.front().score) { + std::pop_heap(heap.begin(), heap.end(), heapCmp); + heap.back() = ScoreIdx{sc, static_cast(i)}; + std::push_heap(heap.begin(), heap.end(), heapCmp); + } + } + std::sort(heap.begin(), heap.end(), [](const ScoreIdx& a, const ScoreIdx& b) { return a.score > b.score; }); + + for (size_t i = 0; i < heap.size(); i++) { + results[i] = Result(heap[i].id, heap[i].score, -1, -1); + } + + Stats stats; + stats.results = heap.size(); + stats.numOverflows = 0; + stats.seconds = seconds; + stats.gcups = seconds > 0 ? (double)ctx->pendingTotalCells / seconds / 1e9 : 0; + return stats; +} + +Marv::Stats Marv::scan(const char* sequence, size_t sequenceLength, int8_t* pssm, Result* results) { + submitScan(sequence, sequenceLength, pssm); + return collectScan(results); +} diff --git a/lib/libmarv/src/marv.h b/lib/libmarv/src/marv.h index 4fcd69749..6975df185 100644 --- a/lib/libmarv/src/marv.h +++ b/lib/libmarv/src/marv.h @@ -47,6 +47,16 @@ class Marv { //sequence must be encoded Stats scan(const char* sequence, size_t sequenceLength, int8_t* pssm, Result* results); + // Optional asynchronous 2-phase form of scan(), letting a caller overlap + // building the next query's profile with this query's GPU execution. + // submitScan() dispatches GPU work without blocking; it must be followed + // by exactly one collectScan() (which blocks until that work completes) + // before submitScan() may be called again. Only the Metal/MPS backend + // implements these; the CUDA/HIP backend does not define them, so + // callers must guard use behind #ifdef HAVE_MPS. + void submitScan(const char* sequence, size_t sequenceLength, int8_t* pssm); + Stats collectScan(Result* results); + private: size_t dbEntries; int alphabetSize; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8dea648ef..dd70fa287 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -284,6 +284,11 @@ if (ENABLE_CUDA) endif () endif () +if (ENABLE_MPS) + target_compile_definitions(mmseqs-framework PUBLIC -DHAVE_MPS=1) + target_link_libraries(mmseqs-framework marv) +endif () + if (NOT FRAMEWORK_ONLY) include(MMseqsSetupDerivedTarget) add_subdirectory(version) diff --git a/src/commons/GpuUtil.cpp b/src/commons/GpuUtil.cpp index 214cd08bc..54c096106 100644 --- a/src/commons/GpuUtil.cpp +++ b/src/commons/GpuUtil.cpp @@ -1,4 +1,4 @@ -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) #include "GpuUtil.h" #include "Debug.h" diff --git a/src/commons/Parameters.cpp b/src/commons/Parameters.cpp index f1815cbdb..5d99b0113 100644 --- a/src/commons/Parameters.cpp +++ b/src/commons/Parameters.cpp @@ -2625,7 +2625,7 @@ void Parameters::setDefaults() { // gpu gpu = 0; -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) char* gpuEnv = getenv("MMSEQS_FORCE_GPU"); if (gpuEnv != NULL) { gpu = 1; @@ -2633,7 +2633,7 @@ void Parameters::setDefaults() { #endif gpuServer = 0; gpuServerWaitTimeout = 10 * 60; -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) char* gpuServerEnv = getenv("MMSEQS_FORCE_GPUSERVER"); if (gpuServerEnv != NULL) { gpuServer = 1; diff --git a/src/prefiltering/ungappedprefilter.cpp b/src/prefiltering/ungappedprefilter.cpp index b03a0f5a6..104386a19 100644 --- a/src/prefiltering/ungappedprefilter.cpp +++ b/src/prefiltering/ungappedprefilter.cpp @@ -25,13 +25,13 @@ #include #endif // #define HAVE_CUDA 1 -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) #include "GpuUtil.h" #include "Alignment.h" #include #endif -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) volatile sig_atomic_t keepRunningClient = 1; void intHandlerClient(int) { @@ -168,6 +168,95 @@ void runFilterOnGpu(Parameters & par, BaseMatrix * subMat, sigaction(SIGTERM, &act, NULL); } + // Processes one query's hits (already sitting in `results`) and writes + // them out. Factored out so the MPS pipelined loop below can call it for + // the *previous* query while the current query's GPU work is in flight. + auto processResults = [&](const Marv::Stats& stats, size_t queryKeyForResults, int queryLenForResults) { + for(size_t i = 0; i < stats.results; i++){ + DBKeyType targetKey = tdbr->getDbKey(results[i].id); + int score = results[i].score; + if(taxonomyHook != NULL){ + TaxID currTax = taxonomyHook->taxonomyMapping->lookup(targetKey); + if (taxonomyHook->expression[0]->isAncestor(currTax) == false) { + continue; + } + } + // check if evalThr != inf + // double evalue = 0.0; + // if (par.evalThr < std::numeric_limits::max()) { + // evalue = evaluer->computeEvalue(score, queryLenForResults); + // } + // bool hasEvalue = (evalue <= par.evalThr); + bool hasDiagScore = (score > par.minDiagScoreThr); + + const bool isIdentity = (queryKeyForResults == targetKey && (par.includeIdentity || sameDB))? true : false; + // --filter-hits + if (isIdentity || hasDiagScore) { + if(par.prefMode == Parameters::PREF_MODE_UNGAPPED_AND_GAPPED){ + Matcher::result_t res; + res.dbKey = targetKey; + res.eval = evaluer->computeEvalue(score, queryLenForResults); + res.dbEndPos = results[i].dbEndPos; + res.dbLen = tdbr->getSeqLen(results[i].id); + res.qEndPos = results[i].qEndPos; + res.qLen = queryLenForResults; + unsigned int qAlnLen = std::max(static_cast(res.qEndPos), static_cast(1)); + unsigned int dbAlnLen = std::max(static_cast(res.dbEndPos), static_cast(1)); + //seqId = (alignment.score1 / static_cast(std::max(dbAlnLen, qAlnLen))) * 0.1656 + 0.1141; + res.seqId = Matcher::estimateSeqIdByScorePerCol(score, qAlnLen, dbAlnLen); + res.qcov = SmithWaterman::computeCov(0, res.qEndPos, res.qLen ); + res.dbcov = SmithWaterman::computeCov(0, res.dbEndPos, res.dbLen ); + res.score = evaluer->computeBitScore(score); + if(Alignment::checkCriteria(res, isIdentity, par.evalThr, par.seqIdThr, par.alnLenThr, par.covMode, par.covThr)){ + resultsAln.emplace_back(res); + } + } else { + hit_t hit; + hit.seqId = targetKey; + hit.prefScore = score; + hit.diagonal = 0; + shortResults.emplace_back(hit); + } + } + } + if(par.prefMode == Parameters::PREF_MODE_UNGAPPED_AND_GAPPED) { + SORT_PARALLEL(resultsAln.begin(), resultsAln.end(), Matcher::compareHits); + size_t maxSeqs = std::min(par.maxResListLen, resultsAln.size()); + for (size_t i = 0; i < maxSeqs; ++i) { + size_t len = Matcher::resultToBuffer(buffer, resultsAln[i], false); + resultBuffer.append(buffer, len); + } + }else{ + SORT_PARALLEL(shortResults.begin(), shortResults.end(), hit_t::compareHitsByScoreAndId); + size_t maxSeqs = std::min(par.maxResListLen, shortResults.size()); + for (size_t i = 0; i < maxSeqs; ++i) { + size_t len = QueryMatcher::prefilterHitToBuffer(buffer, shortResults[i]); + resultBuffer.append(buffer, len); + } + } + + resultWriter.writeData(resultBuffer.c_str(), resultBuffer.length(), queryKeyForResults, 0); + resultBuffer.clear(); + shortResults.clear(); + resultsAln.clear(); + progress.updateProgress(); + }; + + // On the Metal/MPS backend (non-server mode only), pipeline queries: + // submit query i's GPU work without blocking, then build query i+1's + // profile on the CPU while it runs, and only then collect query i's + // results. This overlaps CPU-side profile/composition-bias work with + // GPU execution instead of serializing them. The CUDA/HIP path and + // gpuserver shared-memory protocol are untouched. +#ifdef HAVE_MPS + const bool pipeline = (serverMode == 0); +#else + const bool pipeline = false; +#endif + size_t prevQueryKey = 0; + int prevQueryLen = 0; + bool havePending = false; + // marv.prefetch(); for (size_t id = 0; id < qdbr->getSize(); id++) { if (!keepRunningClient) { @@ -202,6 +291,24 @@ void runFilterOnGpu(Parameters & par, BaseMatrix * subMat, } } } + +#ifdef HAVE_MPS + if (pipeline) { + if (havePending) { + Marv::Stats stats = marv->collectScan(results.data()); + if (keepRunningClient == false) { + EXIT(EXIT_FAILURE); + } + processResults(stats, prevQueryKey, prevQueryLen); + } + marv->submitScan(reinterpret_cast(qSeq.numSequence), qSeq.L, profile); + prevQueryKey = queryKey; + prevQueryLen = qSeq.L; + havePending = true; + continue; + } +#endif + Marv::Stats stats; if (serverMode == 0) { stats = marv->scan(reinterpret_cast(qSeq.numSequence), qSeq.L, profile, results.data()); @@ -259,75 +366,17 @@ void runFilterOnGpu(Parameters & par, BaseMatrix * subMat, EXIT(EXIT_FAILURE); } - for(size_t i = 0; i < stats.results; i++){ - DBKeyType targetKey = tdbr->getDbKey(results[i].id); - int score = results[i].score; - if(taxonomyHook != NULL){ - TaxID currTax = taxonomyHook->taxonomyMapping->lookup(targetKey); - if (taxonomyHook->expression[0]->isAncestor(currTax) == false) { - continue; - } - } - // check if evalThr != inf - // double evalue = 0.0; - // if (par.evalThr < std::numeric_limits::max()) { - // evalue = evaluer->computeEvalue(score, qSeq.L); - // } - // bool hasEvalue = (evalue <= par.evalThr); - bool hasDiagScore = (score > par.minDiagScoreThr); - - const bool isIdentity = (queryKey == targetKey && (par.includeIdentity || sameDB))? true : false; - // --filter-hits - if (isIdentity || hasDiagScore) { - if(par.prefMode == Parameters::PREF_MODE_UNGAPPED_AND_GAPPED){ - Matcher::result_t res; - res.dbKey = targetKey; - res.eval = evaluer->computeEvalue(score, qSeq.L); - res.dbEndPos = results[i].dbEndPos; - res.dbLen = tdbr->getSeqLen(results[i].id); - res.qEndPos = results[i].qEndPos; - res.qLen = qSeq.L; - unsigned int qAlnLen = std::max(static_cast(res.qEndPos), static_cast(1)); - unsigned int dbAlnLen = std::max(static_cast(res.dbEndPos), static_cast(1)); - //seqId = (alignment.score1 / static_cast(std::max(dbAlnLen, qAlnLen))) * 0.1656 + 0.1141; - res.seqId = Matcher::estimateSeqIdByScorePerCol(score, qAlnLen, dbAlnLen); - res.qcov = SmithWaterman::computeCov(0, res.qEndPos, res.qLen ); - res.dbcov = SmithWaterman::computeCov(0, res.dbEndPos, res.dbLen ); - res.score = evaluer->computeBitScore(score); - if(Alignment::checkCriteria(res, isIdentity, par.evalThr, par.seqIdThr, par.alnLenThr, par.covMode, par.covThr)){ - resultsAln.emplace_back(res); - } - } else { - hit_t hit; - hit.seqId = targetKey; - hit.prefScore = score; - hit.diagonal = 0; - shortResults.emplace_back(hit); - } - } - } - if(par.prefMode == Parameters::PREF_MODE_UNGAPPED_AND_GAPPED) { - SORT_PARALLEL(resultsAln.begin(), resultsAln.end(), Matcher::compareHits); - size_t maxSeqs = std::min(par.maxResListLen, resultsAln.size()); - for (size_t i = 0; i < maxSeqs; ++i) { - size_t len = Matcher::resultToBuffer(buffer, resultsAln[i], false); - resultBuffer.append(buffer, len); - } - }else{ - SORT_PARALLEL(shortResults.begin(), shortResults.end(), hit_t::compareHitsByScoreAndId); - size_t maxSeqs = std::min(par.maxResListLen, shortResults.size()); - for (size_t i = 0; i < maxSeqs; ++i) { - size_t len = QueryMatcher::prefilterHitToBuffer(buffer, shortResults[i]); - resultBuffer.append(buffer, len); - } + processResults(stats, queryKey, qSeq.L); + } +#ifdef HAVE_MPS + if (pipeline && havePending) { + Marv::Stats stats = marv->collectScan(results.data()); + if (keepRunningClient == false) { + EXIT(EXIT_FAILURE); } - - resultWriter.writeData(resultBuffer.c_str(), resultBuffer.length(), queryKey, 0); - resultBuffer.clear(); - shortResults.clear(); - resultsAln.clear(); - progress.updateProgress(); + processResults(stats, prevQueryKey, prevQueryLen); } +#endif if (marv != NULL) { delete marv; } else { @@ -554,11 +603,11 @@ int prefilterInternal(int argc, const char **argv, const Command &command, int m taxonomyHook = new QueryMatcherTaxonomyHook(par.db2, tdbr, par.taxonList, par.threads); } if(par.gpu){ -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) runFilterOnGpu(par, subMat, qdbr, tdbr, sameDB, resultWriter, evaluer, taxonomyHook); #else - Debug(Debug::ERROR) << "MMseqs2 was compiled without CUDA support\n"; + Debug(Debug::ERROR) << "MMseqs2 was compiled without GPU support (CUDA/HIP/MPS)\n"; EXIT(EXIT_FAILURE); #endif }else{ diff --git a/src/util/gpuserver.cpp b/src/util/gpuserver.cpp index 84b8239f2..5aef48b74 100644 --- a/src/util/gpuserver.cpp +++ b/src/util/gpuserver.cpp @@ -7,7 +7,7 @@ #include "NucleotideMatrix.h" // #define HAVE_CUDA 1 -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) #include "GpuUtil.h" #endif @@ -24,7 +24,7 @@ void intHandler(int) { int gpuserver(int argc, const char **argv, const Command& command) { Parameters& par = Parameters::getInstance(); par.parseParameters(argc, argv, command, true, 0, 0); -#ifdef HAVE_CUDA +#if defined(HAVE_CUDA) || defined(HAVE_MPS) bool touch = (par.preloadMode != Parameters::PRELOAD_MODE_MMAP); IndexReader dbrIdx(par.db1, par.threads, IndexReader::SEQUENCES, (touch) ? (IndexReader::PRELOAD_INDEX | IndexReader::PRELOAD_DATA) : 0 ); DBReader* dbr = dbrIdx.sequenceReader; From 315df56acc4aa96741bcf05c815b128d104c0079 Mon Sep 17 00:00:00 2001 From: fnachon Date: Thu, 16 Jul 2026 21:37:20 +0200 Subject: [PATCH 2/3] Fix -Werror=unused-variable on CUDA/HIP CI build pipeline/prevQueryKey/prevQueryLen/havePending were declared unconditionally (with pipeline forced to false via an #else branch for non-MPS builds), but only ever referenced from code gated behind #ifdef HAVE_MPS. Under a CUDA/HIP build, that left them declared-but-unused, which upstream's CI build (-Wall -Wextra -Werror) correctly rejects. Move all four declarations inside the same #ifdef HAVE_MPS block as their uses, so they don't exist at all in a CUDA/HIP translation unit. Verified on real hardware (2x RTX A6000, CUDA 13.3, gcc 12.2), not just inferred: built from scratch with -DENABLE_CUDA=1 on the pre-fix commit to reproduce the exact CI warnings, then rebuilt after this fix with the CI's literal compile flags (-Wall -Wextra -Wdisabled-optimization -Werror -pedantic -fno-exceptions) -- clean compile, full link, working binary. Co-Authored-By: Claude Sonnet 5 --- src/prefiltering/ungappedprefilter.cpp | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/prefiltering/ungappedprefilter.cpp b/src/prefiltering/ungappedprefilter.cpp index 104386a19..a30dc08b0 100644 --- a/src/prefiltering/ungappedprefilter.cpp +++ b/src/prefiltering/ungappedprefilter.cpp @@ -247,15 +247,16 @@ void runFilterOnGpu(Parameters & par, BaseMatrix * subMat, // profile on the CPU while it runs, and only then collect query i's // results. This overlaps CPU-side profile/composition-bias work with // GPU execution instead of serializing them. The CUDA/HIP path and - // gpuserver shared-memory protocol are untouched. + // gpuserver shared-memory protocol are untouched. These are only + // referenced from HAVE_MPS-guarded code below, so they're declared + // here under the same guard to avoid an unused-variable error on the + // CUDA/HIP build (which compiles with -Werror). #ifdef HAVE_MPS const bool pipeline = (serverMode == 0); -#else - const bool pipeline = false; -#endif size_t prevQueryKey = 0; int prevQueryLen = 0; bool havePending = false; +#endif // marv.prefetch(); for (size_t id = 0; id < qdbr->getSize(); id++) { From 6fb17ca331bf3d5f8724b0a66feb2a51c230314b Mon Sep 17 00:00:00 2001 From: fnachon Date: Thu, 16 Jul 2026 23:01:41 +0200 Subject: [PATCH 3/3] Fix incorrect masked-residue handling in the gapless kernel Subject bytes were folded via `& 31` to collapse MMseqs2's soft-masked (lowercase, base+32) residue encoding back to the unmasked base residue, modeled on convert.cuh's CaseSensitive_to_CaseInsensitive from libmarv's CUDA code. That turned out to be the wrong model for this kernel: the CPU prefilter path (StripedSmithWaterman.cpp) and libmarv's real CUDA kernels both treat masked/out-of-alphabet residues as suppression -- clamped to the last alphabet index ("X") -- not as a reversible case encoding to fold away. Folding instead of clamping silently un-masks masked positions, scoring them at full value instead of down-weighting them, defeating the purpose of masking (suppressing spurious hits to low-complexity/repeat regions). Fix: clamp out-of-alphabet bytes to alphabetSize-1 once in loadDb(), when the database is loaded, rather than folding per-residue in the kernel's hot loop -- cheaper (one pass over the DB instead of every kernel access) and matches the CPU path's semantics exactly. Found by direct comparison against an independent Metal implementation (github.com/milot-mirdita/mmseqs2, branch metal-gpu-backend) that scores real database sequences correctly. On a real query/target pair with masked residues, this bug had been inflating the ungapped score by ~15% (7124 vs. the correct 6163, independently confirmed via a from-scratch Python DP reference and cross-checked byte-for-byte identical PSSM/subject input data between both implementations) -- not a rare edge case, since any subject with soft-masked residues was affected, worst for near-duplicate/ self-hits and low-complexity regions. Re-verified against the full boundary-length suite and the 500-query/20k-sequence benchmark; cutoff scores and ~95% of top-300 hit sets now match the independent implementation exactly (remaining differences are the same benign tie-break-at-cutoff pattern already documented elsewhere in this file). Co-Authored-By: Claude Sonnet 5 --- lib/libmarv-mps/marv_mps.mm | 43 ++++++++++++++++++++++++++++++++----- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/lib/libmarv-mps/marv_mps.mm b/lib/libmarv-mps/marv_mps.mm index ab10e00af..fdc75ff90 100644 --- a/lib/libmarv-mps/marv_mps.mm +++ b/lib/libmarv-mps/marv_mps.mm @@ -11,10 +11,24 @@ // // Algorithm: H[i][j] = max(0, H[i-1][j-1] + score(query[i], subject[j])), // boundary zero -- the same no-gap recurrence as libmarv's GAPLESS kernel -// (pssmkernels_gapless.cuh). Subject bytes are folded with `& 31` to -// collapse MMseqs2's soft-masked (lowercase, +32) residue encoding back to -// the base alphabet, mirroring libmarv's CaseSensitive_to_CaseInsensitive -// conversion (see convert.cuh). +// (pssmkernels_gapless.cuh). Subject bytes outside the alphabet (masked -- +// makepaddedseqdb encodes soft-masked/lowercase residues as base+32 -- or +// otherwise unrecognized) are clamped to the last valid alphabet index +// ("X") once at DB-load time in loadDb(), matching how both the CPU +// prefilter path and libmarv's real CUDA kernels treat masking: as +// suppression (scored as X), not as a reversible case-encoding to fold +// away. An earlier version of this file instead folded masked bytes back +// to their unmasked identity via `& 31` (modeled on convert.cuh's +// CaseSensitive_to_CaseInsensitive, which turned out not to be what the +// gapless kernel's masking actually means) -- silently unmasking masked +// residues instead of down-weighting them, which is the opposite of +// masking's purpose. Confirmed and fixed by direct comparison against an +// independent Metal implementation (github.com/milot-mirdita/mmseqs2, +// branch metal-gpu-backend) that scores real database sequences correctly: +// on a shared real query, this bug had been silently inflating ungapped +// scores for any subject with masked residues (self-hits, near-duplicates, +// and low-complexity regions worst affected -- observed ~15% too high on +// one real self-alignment), not merely a rare edge case. // // Parallelism (v2, simdgroup-cooperative): a first version of this kernel // used one GPU thread per subject sequence, with its O(queryLen) DP state in @@ -182,7 +196,7 @@ kernel void gaplessAlignCooperative( int32_t best = 0; for (uint32_t j = 0; j < len; j++) { - const uint8_t s = dbData[start + j] & 31; + const uint8_t s = dbData[start + j]; // pre-clamped to alphabet range in loadDb() device const int8_t* row = pssm + (size_t)s * paddedQueryLen; // The only cross-lane dependency: H[i-1][j-1] for this lane's @@ -330,6 +344,25 @@ kernel void gaplessAlignCooperative( throw std::runtime_error("Marv (MPS): failed to allocate GPU buffer for database"); } + // makepaddedseqdb encodes soft-masked (tantan-repeat-masked, lowercase) + // residues as base_value + 32, matching MMseqs2's own aa2num convention. + // The CPU prefilter path (StripedSmithWaterman/ungappedprefilter.cpp) + // treats any masked or otherwise out-of-alphabet byte as "X" -- clamped + // to the last valid alphabet index -- rather than unmasking it back to + // its original residue. Clamping (not folding via e.g. `& 31`) is what + // makes masking actually suppress scores in masked/low-complexity + // regions; done once here at load time (not per kernel access) since it + // depends only on alphabetSize, not on the query. + { + uint8_t* bytes = static_cast([db->data contents]); + const uint8_t maxValid = static_cast(alphabetSize - 1); + for (size_t i = 0; i < dbByteSize; i++) { + if (bytes[i] > maxValid) { + bytes[i] = maxValid; + } + } + } + std::vector offsets64(dbEntries); std::vector lengths32(dbEntries); for (size_t i = 0; i < dbEntries; i++) {