From 2565c8eb292bb189d96d130b8adaa71592ea451e Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 11:17:30 +0000 Subject: [PATCH 01/23] Split basecalling into its own library --- CMakeLists.txt | 29 +-------- cmake/Metal.cmake | 2 +- dorado/basecall/CMakeLists.txt | 62 ++++++++++++++++++++ dorado/{nn => basecall}/CRFModel.cpp | 2 +- dorado/{nn => basecall}/CRFModel.h | 0 dorado/{nn => basecall}/CRFModelConfig.cpp | 0 dorado/{nn => basecall}/CRFModelConfig.h | 0 dorado/{nn => basecall}/CudaCRFModel.cpp | 0 dorado/{nn => basecall}/CudaCRFModel.h | 2 +- dorado/{nn => basecall}/MetalCRFModel.cpp | 0 dorado/{nn => basecall}/MetalCRFModel.h | 0 dorado/{nn => basecall}/ModelRunner.h | 0 dorado/{ => basecall}/decode/CPUDecoder.cpp | 0 dorado/{ => basecall}/decode/CPUDecoder.h | 0 dorado/{ => basecall}/decode/Decoder.h | 0 dorado/{ => basecall}/decode/GPUDecoder.cpp | 4 +- dorado/{ => basecall}/decode/GPUDecoder.h | 0 dorado/{ => basecall}/decode/beam_search.cpp | 0 dorado/{ => basecall}/decode/beam_search.h | 0 dorado/{nn => basecall}/metal/nn.metal | 0 dorado/cli/basecaller.cpp | 2 +- dorado/cli/duplex.cpp | 2 +- dorado/nn/Runners.cpp | 10 ++-- dorado/read_pipeline/BasecallerNode.cpp | 2 +- dorado/read_pipeline/BasecallerNode.h | 2 +- dorado/read_pipeline/Pipelines.cpp | 4 +- dorado/read_pipeline/ScalerNode.cpp | 2 +- dorado/read_pipeline/ScalerNode.h | 2 +- tests/CMakeLists.txt | 1 + tests/CRFModelConfigTest.cpp | 2 +- tests/NodeSmokeTest.cpp | 12 ++-- 31 files changed, 90 insertions(+), 52 deletions(-) create mode 100644 dorado/basecall/CMakeLists.txt rename dorado/{nn => basecall}/CRFModel.cpp (99%) rename dorado/{nn => basecall}/CRFModel.h (100%) rename dorado/{nn => basecall}/CRFModelConfig.cpp (100%) rename dorado/{nn => basecall}/CRFModelConfig.h (100%) rename dorado/{nn => basecall}/CudaCRFModel.cpp (100%) rename dorado/{nn => basecall}/CudaCRFModel.h (98%) rename dorado/{nn => basecall}/MetalCRFModel.cpp (100%) rename dorado/{nn => basecall}/MetalCRFModel.h (100%) rename dorado/{nn => basecall}/ModelRunner.h (100%) rename dorado/{ => basecall}/decode/CPUDecoder.cpp (100%) rename dorado/{ => basecall}/decode/CPUDecoder.h (100%) rename dorado/{ => basecall}/decode/Decoder.h (100%) rename dorado/{ => basecall}/decode/GPUDecoder.cpp (98%) rename dorado/{ => basecall}/decode/GPUDecoder.h (100%) rename dorado/{ => basecall}/decode/beam_search.cpp (100%) rename dorado/{ => basecall}/decode/beam_search.h (100%) rename dorado/{nn => basecall}/metal/nn.metal (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 081457de..ace7db55 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -161,6 +161,7 @@ enable_testing() add_subdirectory(dorado/utils) add_subdirectory(dorado/models) +add_subdirectory(dorado/basecall) if("${CUDA_VERSION}" STREQUAL "") set(CUDA_VERSION ${CUDAToolkit_VERSION}) @@ -179,11 +180,6 @@ set(LIB_SOURCE_FILES dorado/alignment/Minimap2Index.h dorado/alignment/Minimap2IndexSupportTypes.h dorado/alignment/Minimap2Options.h - dorado/nn/CRFModel.h - dorado/nn/CRFModel.cpp - dorado/nn/CRFModelConfig.h - dorado/nn/CRFModelConfig.cpp - dorado/nn/ModelRunner.h dorado/nn/ModBaseModel.cpp dorado/nn/ModBaseModel.h dorado/nn/ModBaseModelConfig.cpp @@ -263,10 +259,6 @@ set(LIB_SOURCE_FILES dorado/demux/Trimmer.h dorado/demux/parse_custom_kit.cpp dorado/demux/parse_custom_kit.h - dorado/decode/beam_search.cpp - dorado/decode/beam_search.h - dorado/decode/CPUDecoder.cpp - dorado/decode/CPUDecoder.h dorado/modbase/modbase_encoder.cpp dorado/modbase/modbase_encoder.h dorado/modbase/modbase_scaler.cpp @@ -277,22 +269,6 @@ set(LIB_SOURCE_FILES dorado/modbase/MotifMatcher.h ) -if (DORADO_GPU_BUILD) - if(APPLE) - list(APPEND LIB_SOURCE_FILES - dorado/nn/MetalCRFModel.h - dorado/nn/MetalCRFModel.cpp - ) - else() - list(APPEND LIB_SOURCE_FILES - dorado/decode/GPUDecoder.cpp - dorado/decode/GPUDecoder.h - dorado/nn/CudaCRFModel.h - dorado/nn/CudaCRFModel.cpp - ) - endif() -endif() - add_library(dorado_lib ${LIB_SOURCE_FILES}) enable_warnings_as_errors(dorado_lib) @@ -312,7 +288,6 @@ target_include_directories(dorado_lib PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/dorado ${CMAKE_CURRENT_SOURCE_DIR}/dorado - ${KOI_INCLUDE} ${POD5_INCLUDE} ) # 3rdparty libs should be considered SYSTEM headers @@ -348,8 +323,8 @@ target_link_libraries(dorado_lib vbz_hdf_plugin edlib dorado_utils + dorado_basecall PRIVATE - ${KOI_LIBRARIES} minimap2 ) diff --git a/cmake/Metal.cmake b/cmake/Metal.cmake index caf02e91..212fe673 100644 --- a/cmake/Metal.cmake +++ b/cmake/Metal.cmake @@ -5,7 +5,7 @@ find_library(APPLE_FWK_METAL Metal REQUIRED) find_library(IOKIT IOKit REQUIRED) set(AIR_FILES) -set(METAL_SOURCES dorado/nn/metal/nn.metal) +set(METAL_SOURCES dorado/basecall/metal/nn.metal) if (CMAKE_SYSTEM_NAME STREQUAL "iOS") set(XCRUN_SDK iphoneos) diff --git a/dorado/basecall/CMakeLists.txt b/dorado/basecall/CMakeLists.txt new file mode 100644 index 00000000..d842b484 --- /dev/null +++ b/dorado/basecall/CMakeLists.txt @@ -0,0 +1,62 @@ + + +set(LIB_SOURCE_FILES + CRFModel.h + CRFModel.cpp + CRFModelConfig.h + CRFModelConfig.cpp + ModelRunner.h + decode/beam_search.cpp + decode/beam_search.h + decode/CPUDecoder.cpp + decode/CPUDecoder.h +) + +if (DORADO_GPU_BUILD) + if(APPLE) + list(APPEND LIB_SOURCE_FILES + MetalCRFModel.h + MetalCRFModel.cpp + ) + else() + list(APPEND LIB_SOURCE_FILES + CudaCRFModel.h + CudaCRFModel.cpp + decode/GPUDecoder.cpp + decode/GPUDecoder.h + ) + endif() +endif() + +add_library(dorado_basecall STATIC ${LIB_SOURCE_FILES}) + +target_compile_definitions(dorado_basecall + PUBLIC + DORADO_GPU_BUILD=$ +) + +target_include_directories(dorado_basecall + SYSTEM + PUBLIC + ${DORADO_3RD_PARTY_SOURCE}/toml11 + PRIVATE + ${DORADO_3RD_PARTY_SOURCE}/NVTX/c/include +) + + +target_link_libraries(dorado_basecall + PUBLIC + ${TORCH_LIBRARIES} + dorado_utils + dorado_models_lib + PRIVATE + ${KOI_LIBRARIES} + spdlog::spdlog +) + +target_include_directories(dorado_basecall + PRIVATE + ${KOI_INCLUDE} +) + +enable_warnings_as_errors(dorado_basecall) diff --git a/dorado/nn/CRFModel.cpp b/dorado/basecall/CRFModel.cpp similarity index 99% rename from dorado/nn/CRFModel.cpp rename to dorado/basecall/CRFModel.cpp index 2dd7db46..7ba29561 100644 --- a/dorado/nn/CRFModel.cpp +++ b/dorado/basecall/CRFModel.cpp @@ -12,7 +12,7 @@ #if DORADO_GPU_BUILD && !defined(__APPLE__) #define USE_KOI 1 -#include "../utils/cuda_utils.h" +#include "utils/cuda_utils.h" #include #include diff --git a/dorado/nn/CRFModel.h b/dorado/basecall/CRFModel.h similarity index 100% rename from dorado/nn/CRFModel.h rename to dorado/basecall/CRFModel.h diff --git a/dorado/nn/CRFModelConfig.cpp b/dorado/basecall/CRFModelConfig.cpp similarity index 100% rename from dorado/nn/CRFModelConfig.cpp rename to dorado/basecall/CRFModelConfig.cpp diff --git a/dorado/nn/CRFModelConfig.h b/dorado/basecall/CRFModelConfig.h similarity index 100% rename from dorado/nn/CRFModelConfig.h rename to dorado/basecall/CRFModelConfig.h diff --git a/dorado/nn/CudaCRFModel.cpp b/dorado/basecall/CudaCRFModel.cpp similarity index 100% rename from dorado/nn/CudaCRFModel.cpp rename to dorado/basecall/CudaCRFModel.cpp diff --git a/dorado/nn/CudaCRFModel.h b/dorado/basecall/CudaCRFModel.h similarity index 98% rename from dorado/nn/CudaCRFModel.h rename to dorado/basecall/CudaCRFModel.h index c216164f..77720cdf 100644 --- a/dorado/nn/CudaCRFModel.h +++ b/dorado/basecall/CudaCRFModel.h @@ -1,7 +1,7 @@ #pragma once +#include "CRFModel.h" #include "ModelRunner.h" -#include "nn/CRFModel.h" #include #include diff --git a/dorado/nn/MetalCRFModel.cpp b/dorado/basecall/MetalCRFModel.cpp similarity index 100% rename from dorado/nn/MetalCRFModel.cpp rename to dorado/basecall/MetalCRFModel.cpp diff --git a/dorado/nn/MetalCRFModel.h b/dorado/basecall/MetalCRFModel.h similarity index 100% rename from dorado/nn/MetalCRFModel.h rename to dorado/basecall/MetalCRFModel.h diff --git a/dorado/nn/ModelRunner.h b/dorado/basecall/ModelRunner.h similarity index 100% rename from dorado/nn/ModelRunner.h rename to dorado/basecall/ModelRunner.h diff --git a/dorado/decode/CPUDecoder.cpp b/dorado/basecall/decode/CPUDecoder.cpp similarity index 100% rename from dorado/decode/CPUDecoder.cpp rename to dorado/basecall/decode/CPUDecoder.cpp diff --git a/dorado/decode/CPUDecoder.h b/dorado/basecall/decode/CPUDecoder.h similarity index 100% rename from dorado/decode/CPUDecoder.h rename to dorado/basecall/decode/CPUDecoder.h diff --git a/dorado/decode/Decoder.h b/dorado/basecall/decode/Decoder.h similarity index 100% rename from dorado/decode/Decoder.h rename to dorado/basecall/decode/Decoder.h diff --git a/dorado/decode/GPUDecoder.cpp b/dorado/basecall/decode/GPUDecoder.cpp similarity index 98% rename from dorado/decode/GPUDecoder.cpp rename to dorado/basecall/decode/GPUDecoder.cpp index 7718f3f0..f3fa7766 100644 --- a/dorado/decode/GPUDecoder.cpp +++ b/dorado/basecall/decode/GPUDecoder.cpp @@ -1,8 +1,8 @@ #include "GPUDecoder.h" -#include "../utils/cuda_utils.h" -#include "../utils/gpu_profiling.h" #include "Decoder.h" +#include "utils/cuda_utils.h" +#include "utils/gpu_profiling.h" #include #include diff --git a/dorado/decode/GPUDecoder.h b/dorado/basecall/decode/GPUDecoder.h similarity index 100% rename from dorado/decode/GPUDecoder.h rename to dorado/basecall/decode/GPUDecoder.h diff --git a/dorado/decode/beam_search.cpp b/dorado/basecall/decode/beam_search.cpp similarity index 100% rename from dorado/decode/beam_search.cpp rename to dorado/basecall/decode/beam_search.cpp diff --git a/dorado/decode/beam_search.h b/dorado/basecall/decode/beam_search.h similarity index 100% rename from dorado/decode/beam_search.h rename to dorado/basecall/decode/beam_search.h diff --git a/dorado/nn/metal/nn.metal b/dorado/basecall/metal/nn.metal similarity index 100% rename from dorado/nn/metal/nn.metal rename to dorado/basecall/metal/nn.metal diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp index 4603b4e8..192ce3b2 100644 --- a/dorado/cli/basecaller.cpp +++ b/dorado/cli/basecaller.cpp @@ -1,10 +1,10 @@ #include "Version.h" +#include "basecall/CRFModelConfig.h" #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" #include "models/kits.h" #include "models/models.h" -#include "nn/CRFModelConfig.h" #include "nn/ModBaseRunner.h" #include "nn/Runners.h" #include "read_pipeline/AdapterDetectorNode.h" diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index ffd14ad4..c9803822 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -1,11 +1,11 @@ #include "Version.h" +#include "basecall/CRFModelConfig.h" #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" #include "models/kits.h" #include "models/metadata.h" #include "models/models.h" -#include "nn/CRFModelConfig.h" #include "nn/ModBaseRunner.h" #include "nn/Runners.h" #include "read_pipeline/AlignerNode.h" diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index c4787fe0..8cd13bd4 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -1,15 +1,15 @@ #include "Runners.h" #include "ModBaseRunner.h" -#include "ModelRunner.h" -#include "decode/CPUDecoder.h" -#include "nn/CRFModel.h" +#include "basecall/CRFModel.h" +#include "basecall/ModelRunner.h" +#include "basecall/decode/CPUDecoder.h" #if DORADO_GPU_BUILD #ifdef __APPLE__ -#include "nn/MetalCRFModel.h" +#include "basecall/MetalCRFModel.h" #else -#include "nn/CudaCRFModel.h" +#include "basecall/CudaCRFModel.h" #include "utils/cuda_utils.h" #endif #endif // DORADO_GPU_BUILD diff --git a/dorado/read_pipeline/BasecallerNode.cpp b/dorado/read_pipeline/BasecallerNode.cpp index e29bcaed..7048029c 100644 --- a/dorado/read_pipeline/BasecallerNode.cpp +++ b/dorado/read_pipeline/BasecallerNode.cpp @@ -1,6 +1,6 @@ #include "BasecallerNode.h" -#include "decode/CPUDecoder.h" +#include "basecall/decode/CPUDecoder.h" #include "stitch.h" #include "utils/stats.h" diff --git a/dorado/read_pipeline/BasecallerNode.h b/dorado/read_pipeline/BasecallerNode.h index 9a3236e5..9f0f9862 100644 --- a/dorado/read_pipeline/BasecallerNode.h +++ b/dorado/read_pipeline/BasecallerNode.h @@ -1,7 +1,7 @@ #pragma once -#include "../nn/ModelRunner.h" #include "ReadPipeline.h" +#include "basecall/ModelRunner.h" #include "utils/AsyncQueue.h" #include "utils/stats.h" diff --git a/dorado/read_pipeline/Pipelines.cpp b/dorado/read_pipeline/Pipelines.cpp index d1fb622f..80ed4ed6 100644 --- a/dorado/read_pipeline/Pipelines.cpp +++ b/dorado/read_pipeline/Pipelines.cpp @@ -6,9 +6,9 @@ #include "ReadSplitNode.h" #include "ScalerNode.h" #include "StereoDuplexEncoderNode.h" -#include "nn/CRFModelConfig.h" +#include "basecall/CRFModelConfig.h" +#include "basecall/ModelRunner.h" #include "nn/ModBaseRunner.h" -#include "nn/ModelRunner.h" #include "splitter/DuplexReadSplitter.h" #include "splitter/RNAReadSplitter.h" diff --git a/dorado/read_pipeline/ScalerNode.cpp b/dorado/read_pipeline/ScalerNode.cpp index 6fa77e79..2be915e6 100644 --- a/dorado/read_pipeline/ScalerNode.cpp +++ b/dorado/read_pipeline/ScalerNode.cpp @@ -1,6 +1,6 @@ #include "ScalerNode.h" -#include "nn/CRFModelConfig.h" +#include "basecall/CRFModelConfig.h" #include "utils/tensor_utils.h" #include "utils/trim.h" diff --git a/dorado/read_pipeline/ScalerNode.h b/dorado/read_pipeline/ScalerNode.h index 64f65c11..4bd521a0 100644 --- a/dorado/read_pipeline/ScalerNode.h +++ b/dorado/read_pipeline/ScalerNode.h @@ -1,6 +1,6 @@ #pragma once #include "ReadPipeline.h" -#include "nn/CRFModelConfig.h" +#include "basecall/CRFModelConfig.h" #include "utils/stats.h" #include diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8d34f125..b612e452 100755 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -75,6 +75,7 @@ foreach(TEST_BIN dorado_tests dorado_smoke_tests) dorado_lib dorado_io_lib dorado_models_lib + dorado_basecall minimap2 ${ZLIB_LIBRARIES} ${POD5_LIBRARIES} diff --git a/tests/CRFModelConfigTest.cpp b/tests/CRFModelConfigTest.cpp index e6dc0342..6a47224d 100644 --- a/tests/CRFModelConfigTest.cpp +++ b/tests/CRFModelConfigTest.cpp @@ -1,4 +1,4 @@ -#include "nn/CRFModelConfig.h" +#include "basecall/CRFModelConfig.h" #include "TestUtils.h" diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 8f2cca06..0e1978d2 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -1,12 +1,12 @@ #include "MessageSinkUtils.h" #include "TestUtils.h" -#include "decode/CPUDecoder.h" +#include "basecall/CRFModel.h" +#include "basecall/CRFModelConfig.h" +#include "basecall/ModelRunner.h" +#include "basecall/decode/CPUDecoder.h" #include "models/models.h" -#include "nn/CRFModel.h" -#include "nn/CRFModelConfig.h" #include "nn/ModBaseModel.h" #include "nn/ModBaseRunner.h" -#include "nn/ModelRunner.h" #include "read_pipeline/AdapterDetectorNode.h" #include "read_pipeline/BarcodeClassifierNode.h" #include "read_pipeline/BasecallerNode.h" @@ -21,9 +21,9 @@ #if DORADO_GPU_BUILD #ifdef __APPLE__ -#include "nn/MetalCRFModel.h" +#include "basecall/MetalCRFModel.h" #else -#include "nn/CudaCRFModel.h" +#include "basecall/CudaCRFModel.h" #include "utils/cuda_utils.h" #endif #endif // DORADO_GPU_BUILD From cb585ea5cc3b1aa9016bafaf9aaba537cca6cf4c Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 11:17:59 +0000 Subject: [PATCH 02/23] Split modbase into its own library --- CMakeLists.txt | 16 +------ dorado/basecall/CMakeLists.txt | 2 - dorado/cli/basecaller.cpp | 2 +- dorado/cli/duplex.cpp | 2 +- dorado/modbase/CMakeLists.txt | 42 +++++++++++++++++++ dorado/{nn => modbase}/ModBaseModel.cpp | 0 dorado/{nn => modbase}/ModBaseModel.h | 0 dorado/{nn => modbase}/ModBaseModelConfig.cpp | 0 dorado/{nn => modbase}/ModBaseModelConfig.h | 0 dorado/{nn => modbase}/ModBaseRunner.cpp | 4 +- dorado/{nn => modbase}/ModBaseRunner.h | 0 ...modbase_encoder.cpp => ModbaseEncoder.cpp} | 2 +- .../{modbase_encoder.h => ModbaseEncoder.h} | 0 .../{modbase_scaler.cpp => ModbaseScaler.cpp} | 2 +- .../{modbase_scaler.h => ModbaseScaler.h} | 0 dorado/modbase/MotifMatcher.cpp | 2 +- dorado/nn/Runners.cpp | 2 +- dorado/read_pipeline/ModBaseCallerNode.cpp | 6 +-- dorado/read_pipeline/Pipelines.cpp | 2 +- tests/CMakeLists.txt | 1 + tests/ModBaseEncoderTest.cpp | 3 +- tests/MotifMatcherTest.cpp | 2 +- tests/NodeSmokeTest.cpp | 4 +- 23 files changed, 62 insertions(+), 32 deletions(-) create mode 100644 dorado/modbase/CMakeLists.txt rename dorado/{nn => modbase}/ModBaseModel.cpp (100%) rename dorado/{nn => modbase}/ModBaseModel.h (100%) rename dorado/{nn => modbase}/ModBaseModelConfig.cpp (100%) rename dorado/{nn => modbase}/ModBaseModelConfig.h (100%) rename dorado/{nn => modbase}/ModBaseRunner.cpp (99%) rename dorado/{nn => modbase}/ModBaseRunner.h (100%) rename dorado/modbase/{modbase_encoder.cpp => ModbaseEncoder.cpp} (99%) rename dorado/modbase/{modbase_encoder.h => ModbaseEncoder.h} (100%) rename dorado/modbase/{modbase_scaler.cpp => ModbaseScaler.cpp} (99%) rename dorado/modbase/{modbase_scaler.h => ModbaseScaler.h} (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index ace7db55..e33aa3e2 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -162,6 +162,7 @@ enable_testing() add_subdirectory(dorado/utils) add_subdirectory(dorado/models) add_subdirectory(dorado/basecall) +add_subdirectory(dorado/modbase) if("${CUDA_VERSION}" STREQUAL "") set(CUDA_VERSION ${CUDAToolkit_VERSION}) @@ -180,12 +181,6 @@ set(LIB_SOURCE_FILES dorado/alignment/Minimap2Index.h dorado/alignment/Minimap2IndexSupportTypes.h dorado/alignment/Minimap2Options.h - dorado/nn/ModBaseModel.cpp - dorado/nn/ModBaseModel.h - dorado/nn/ModBaseModelConfig.cpp - dorado/nn/ModBaseModelConfig.h - dorado/nn/ModBaseRunner.cpp - dorado/nn/ModBaseRunner.h dorado/nn/Runners.cpp dorado/nn/Runners.h dorado/read_pipeline/Pipelines.cpp @@ -259,14 +254,6 @@ set(LIB_SOURCE_FILES dorado/demux/Trimmer.h dorado/demux/parse_custom_kit.cpp dorado/demux/parse_custom_kit.h - dorado/modbase/modbase_encoder.cpp - dorado/modbase/modbase_encoder.h - dorado/modbase/modbase_scaler.cpp - dorado/modbase/modbase_scaler.h - dorado/modbase/ModBaseContext.cpp - dorado/modbase/ModBaseContext.h - dorado/modbase/MotifMatcher.cpp - dorado/modbase/MotifMatcher.h ) add_library(dorado_lib ${LIB_SOURCE_FILES}) @@ -324,6 +311,7 @@ target_link_libraries(dorado_lib edlib dorado_utils dorado_basecall + dorado_modbase PRIVATE minimap2 ) diff --git a/dorado/basecall/CMakeLists.txt b/dorado/basecall/CMakeLists.txt index d842b484..3acb7323 100644 --- a/dorado/basecall/CMakeLists.txt +++ b/dorado/basecall/CMakeLists.txt @@ -1,5 +1,3 @@ - - set(LIB_SOURCE_FILES CRFModel.h CRFModel.cpp diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp index 192ce3b2..e4aed1d9 100644 --- a/dorado/cli/basecaller.cpp +++ b/dorado/cli/basecaller.cpp @@ -3,9 +3,9 @@ #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" +#include "modbase/ModBaseRunner.h" #include "models/kits.h" #include "models/models.h" -#include "nn/ModBaseRunner.h" #include "nn/Runners.h" #include "read_pipeline/AdapterDetectorNode.h" #include "read_pipeline/AlignerNode.h" diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index c9803822..edcd6ba5 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -3,10 +3,10 @@ #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" +#include "modbase/ModBaseRunner.h" #include "models/kits.h" #include "models/metadata.h" #include "models/models.h" -#include "nn/ModBaseRunner.h" #include "nn/Runners.h" #include "read_pipeline/AlignerNode.h" #include "read_pipeline/BaseSpaceDuplexCallerNode.h" diff --git a/dorado/modbase/CMakeLists.txt b/dorado/modbase/CMakeLists.txt new file mode 100644 index 00000000..3a5d0695 --- /dev/null +++ b/dorado/modbase/CMakeLists.txt @@ -0,0 +1,42 @@ +set(LIB_SOURCE_FILES + ModBaseContext.cpp + ModBaseContext.h + ModbaseEncoder.cpp + ModbaseEncoder.h + ModBaseModel.cpp + ModBaseModel.h + ModBaseModelConfig.cpp + ModBaseModelConfig.h + ModBaseRunner.cpp + ModBaseRunner.h + ModbaseScaler.cpp + ModbaseScaler.h + MotifMatcher.cpp + MotifMatcher.h +) + +add_library(dorado_modbase STATIC ${LIB_SOURCE_FILES}) + +target_compile_definitions(dorado_modbase + PUBLIC + DORADO_GPU_BUILD=$ +) + +target_include_directories(dorado_modbase + SYSTEM + PUBLIC + PRIVATE + ${DORADO_3RD_PARTY_SOURCE}/NVTX/c/include + ${DORADO_3RD_PARTY_SOURCE}/toml11 +) + + +target_link_libraries(dorado_modbase + PUBLIC + ${TORCH_LIBRARIES} + PRIVATE + dorado_utils + spdlog::spdlog +) + +enable_warnings_as_errors(dorado_modbase) diff --git a/dorado/nn/ModBaseModel.cpp b/dorado/modbase/ModBaseModel.cpp similarity index 100% rename from dorado/nn/ModBaseModel.cpp rename to dorado/modbase/ModBaseModel.cpp diff --git a/dorado/nn/ModBaseModel.h b/dorado/modbase/ModBaseModel.h similarity index 100% rename from dorado/nn/ModBaseModel.h rename to dorado/modbase/ModBaseModel.h diff --git a/dorado/nn/ModBaseModelConfig.cpp b/dorado/modbase/ModBaseModelConfig.cpp similarity index 100% rename from dorado/nn/ModBaseModelConfig.cpp rename to dorado/modbase/ModBaseModelConfig.cpp diff --git a/dorado/nn/ModBaseModelConfig.h b/dorado/modbase/ModBaseModelConfig.h similarity index 100% rename from dorado/nn/ModBaseModelConfig.h rename to dorado/modbase/ModBaseModelConfig.h diff --git a/dorado/nn/ModBaseRunner.cpp b/dorado/modbase/ModBaseRunner.cpp similarity index 99% rename from dorado/nn/ModBaseRunner.cpp rename to dorado/modbase/ModBaseRunner.cpp index 75139f9d..b55f9bf9 100644 --- a/dorado/nn/ModBaseRunner.cpp +++ b/dorado/modbase/ModBaseRunner.cpp @@ -2,8 +2,8 @@ #include "ModBaseModel.h" #include "ModBaseModelConfig.h" -#include "modbase/MotifMatcher.h" -#include "modbase/modbase_scaler.h" +#include "ModbaseScaler.h" +#include "MotifMatcher.h" #include "utils/sequence_utils.h" #include "utils/stats.h" #include "utils/tensor_utils.h" diff --git a/dorado/nn/ModBaseRunner.h b/dorado/modbase/ModBaseRunner.h similarity index 100% rename from dorado/nn/ModBaseRunner.h rename to dorado/modbase/ModBaseRunner.h diff --git a/dorado/modbase/modbase_encoder.cpp b/dorado/modbase/ModbaseEncoder.cpp similarity index 99% rename from dorado/modbase/modbase_encoder.cpp rename to dorado/modbase/ModbaseEncoder.cpp index 3ccc3b7f..77f6943d 100644 --- a/dorado/modbase/modbase_encoder.cpp +++ b/dorado/modbase/ModbaseEncoder.cpp @@ -1,4 +1,4 @@ -#include "modbase_encoder.h" +#include "ModbaseEncoder.h" #include "utils/sequence_utils.h" #include "utils/simd.h" diff --git a/dorado/modbase/modbase_encoder.h b/dorado/modbase/ModbaseEncoder.h similarity index 100% rename from dorado/modbase/modbase_encoder.h rename to dorado/modbase/ModbaseEncoder.h diff --git a/dorado/modbase/modbase_scaler.cpp b/dorado/modbase/ModbaseScaler.cpp similarity index 99% rename from dorado/modbase/modbase_scaler.cpp rename to dorado/modbase/ModbaseScaler.cpp index 1b5a0d18..7817fede 100644 --- a/dorado/modbase/modbase_scaler.cpp +++ b/dorado/modbase/ModbaseScaler.cpp @@ -1,4 +1,4 @@ -#include "modbase_scaler.h" +#include "ModbaseScaler.h" #include "utils/math_utils.h" diff --git a/dorado/modbase/modbase_scaler.h b/dorado/modbase/ModbaseScaler.h similarity index 100% rename from dorado/modbase/modbase_scaler.h rename to dorado/modbase/ModbaseScaler.h diff --git a/dorado/modbase/MotifMatcher.cpp b/dorado/modbase/MotifMatcher.cpp index 0911c802..4570d939 100644 --- a/dorado/modbase/MotifMatcher.cpp +++ b/dorado/modbase/MotifMatcher.cpp @@ -1,6 +1,6 @@ #include "MotifMatcher.h" -#include "nn/ModBaseModelConfig.h" +#include "ModBaseModelConfig.h" #include diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index 8cd13bd4..f743195e 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -1,9 +1,9 @@ #include "Runners.h" -#include "ModBaseRunner.h" #include "basecall/CRFModel.h" #include "basecall/ModelRunner.h" #include "basecall/decode/CPUDecoder.h" +#include "modbase/ModBaseRunner.h" #if DORADO_GPU_BUILD #ifdef __APPLE__ diff --git a/dorado/read_pipeline/ModBaseCallerNode.cpp b/dorado/read_pipeline/ModBaseCallerNode.cpp index 974d340f..ecdce869 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.cpp +++ b/dorado/read_pipeline/ModBaseCallerNode.cpp @@ -1,9 +1,9 @@ #include "ModBaseCallerNode.h" #include "modbase/ModBaseContext.h" -#include "modbase/modbase_encoder.h" -#include "nn/ModBaseModelConfig.h" -#include "nn/ModBaseRunner.h" +#include "modbase/ModBaseModelConfig.h" +#include "modbase/ModBaseRunner.h" +#include "modbase/ModbaseEncoder.h" #include "utils/math_utils.h" #include "utils/sequence_utils.h" #include "utils/stats.h" diff --git a/dorado/read_pipeline/Pipelines.cpp b/dorado/read_pipeline/Pipelines.cpp index 80ed4ed6..435fa4b2 100644 --- a/dorado/read_pipeline/Pipelines.cpp +++ b/dorado/read_pipeline/Pipelines.cpp @@ -8,7 +8,7 @@ #include "StereoDuplexEncoderNode.h" #include "basecall/CRFModelConfig.h" #include "basecall/ModelRunner.h" -#include "nn/ModBaseRunner.h" +#include "modbase/ModBaseRunner.h" #include "splitter/DuplexReadSplitter.h" #include "splitter/RNAReadSplitter.h" diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b612e452..ec727dc0 100755 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -76,6 +76,7 @@ foreach(TEST_BIN dorado_tests dorado_smoke_tests) dorado_io_lib dorado_models_lib dorado_basecall + dorado_modbase minimap2 ${ZLIB_LIBRARIES} ${POD5_LIBRARIES} diff --git a/tests/ModBaseEncoderTest.cpp b/tests/ModBaseEncoderTest.cpp index e29754e2..6ed0dd07 100644 --- a/tests/ModBaseEncoderTest.cpp +++ b/tests/ModBaseEncoderTest.cpp @@ -1,4 +1,5 @@ -#include "modbase/modbase_encoder.h" +#include "modbase/ModbaseEncoder.h" + #include "utils/sequence_utils.h" #include diff --git a/tests/MotifMatcherTest.cpp b/tests/MotifMatcherTest.cpp index 77e380dc..d62207d7 100644 --- a/tests/MotifMatcherTest.cpp +++ b/tests/MotifMatcherTest.cpp @@ -1,6 +1,6 @@ #include "modbase/MotifMatcher.h" -#include "nn/ModBaseModelConfig.h" +#include "modbase/ModBaseModelConfig.h" #include diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 0e1978d2..881ff9d1 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -4,9 +4,9 @@ #include "basecall/CRFModelConfig.h" #include "basecall/ModelRunner.h" #include "basecall/decode/CPUDecoder.h" +#include "modbase/ModBaseModel.h" +#include "modbase/ModBaseRunner.h" #include "models/models.h" -#include "nn/ModBaseModel.h" -#include "nn/ModBaseRunner.h" #include "read_pipeline/AdapterDetectorNode.h" #include "read_pipeline/BarcodeClassifierNode.h" #include "read_pipeline/BasecallerNode.h" From aa413a395c4029dcd69ca2e30e43120338558342 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 12:04:53 +0000 Subject: [PATCH 03/23] Split some common crf functions from the model --- dorado/basecall/CMakeLists.txt | 2 + dorado/basecall/CRFModel.cpp | 67 +-------------------------- dorado/basecall/CRFModel.h | 11 ----- dorado/basecall/MetalCRFModel.cpp | 1 + dorado/basecall/crf_utils.cpp | 76 +++++++++++++++++++++++++++++++ dorado/basecall/crf_utils.h | 20 ++++++++ dorado/nn/Runners.cpp | 1 + 7 files changed, 101 insertions(+), 77 deletions(-) create mode 100644 dorado/basecall/crf_utils.cpp create mode 100644 dorado/basecall/crf_utils.h diff --git a/dorado/basecall/CMakeLists.txt b/dorado/basecall/CMakeLists.txt index 3acb7323..300fe402 100644 --- a/dorado/basecall/CMakeLists.txt +++ b/dorado/basecall/CMakeLists.txt @@ -1,4 +1,6 @@ set(LIB_SOURCE_FILES + crf_utils.h + crf_utils.cpp CRFModel.h CRFModel.cpp CRFModelConfig.h diff --git a/dorado/basecall/CRFModel.cpp b/dorado/basecall/CRFModel.cpp index 7ba29561..14f6727e 100644 --- a/dorado/basecall/CRFModel.cpp +++ b/dorado/basecall/CRFModel.cpp @@ -1,9 +1,9 @@ #include "CRFModel.h" #include "CRFModelConfig.h" +#include "crf_utils.h" #include "utils/gpu_profiling.h" #include "utils/math_utils.h" -#include "utils/memory_utils.h" #include "utils/module_utils.h" #include "utils/tensor_utils.h" @@ -906,45 +906,6 @@ TORCH_MODULE(CRFModel); } // namespace nn -std::vector load_crf_model_weights(const std::filesystem::path &dir, - bool decomposition, - bool linear_layer_bias) { - auto tensors = std::vector{ - - "0.conv.weight.tensor", "0.conv.bias.tensor", - - "1.conv.weight.tensor", "1.conv.bias.tensor", - - "2.conv.weight.tensor", "2.conv.bias.tensor", - - "4.rnn.weight_ih_l0.tensor", "4.rnn.weight_hh_l0.tensor", - "4.rnn.bias_ih_l0.tensor", "4.rnn.bias_hh_l0.tensor", - - "5.rnn.weight_ih_l0.tensor", "5.rnn.weight_hh_l0.tensor", - "5.rnn.bias_ih_l0.tensor", "5.rnn.bias_hh_l0.tensor", - - "6.rnn.weight_ih_l0.tensor", "6.rnn.weight_hh_l0.tensor", - "6.rnn.bias_ih_l0.tensor", "6.rnn.bias_hh_l0.tensor", - - "7.rnn.weight_ih_l0.tensor", "7.rnn.weight_hh_l0.tensor", - "7.rnn.bias_ih_l0.tensor", "7.rnn.bias_hh_l0.tensor", - - "8.rnn.weight_ih_l0.tensor", "8.rnn.weight_hh_l0.tensor", - "8.rnn.bias_ih_l0.tensor", "8.rnn.bias_hh_l0.tensor", - - "9.linear.weight.tensor"}; - - if (linear_layer_bias) { - tensors.push_back("9.linear.bias.tensor"); - } - - if (decomposition) { - tensors.push_back("10.linear.weight.tensor"); - } - - return utils::load_tensors(dir, tensors); -} - ModuleHolder load_crf_model(const CRFModelConfig &model_config, const at::TensorOptions &options) { auto model = nn::CRFModel(model_config); @@ -952,30 +913,4 @@ ModuleHolder load_crf_model(const CRFModelConfig &model_config, model_config.out_features.has_value(), model_config.bias); } -size_t auto_calculate_num_runners(const CRFModelConfig &model_config, - size_t batch_size, - float memory_fraction) { - auto model_name = std::filesystem::canonical(model_config.model_path).filename().string(); - - // very hand-wavy determination - // these numbers were determined empirically by running 1, 2, 4 and 8 runners for each model - auto required_ram_per_runner_GB = 0.f; - if (model_name.find("_fast@v") != std::string::npos) { - required_ram_per_runner_GB = 1.5; - } else if (model_name.find("_hac@v") != std::string::npos) { - required_ram_per_runner_GB = 4.5; - } else if (model_name.find("_sup@v") != std::string::npos) { - required_ram_per_runner_GB = 12.5; - } else { - return 1; - } - - // numbers were determined with a batch_size of 128, assume this just scales - required_ram_per_runner_GB *= batch_size / 128.f; - - auto free_ram_GB = utils::available_host_memory_GB() * memory_fraction; - auto num_runners = static_cast(free_ram_GB / required_ram_per_runner_GB); - return std::clamp(num_runners, size_t(1), std::size_t(std::thread::hardware_concurrency())); -} - } // namespace dorado diff --git a/dorado/basecall/CRFModel.h b/dorado/basecall/CRFModel.h index 03bddd75..6cc9cb58 100644 --- a/dorado/basecall/CRFModel.h +++ b/dorado/basecall/CRFModel.h @@ -2,22 +2,11 @@ #include -#include -#include - namespace dorado { struct CRFModelConfig; -std::vector load_crf_model_weights(const std::filesystem::path& dir, - bool decomposition, - bool bias); - torch::nn::ModuleHolder load_crf_model(const CRFModelConfig& model_config, const at::TensorOptions& options); -size_t auto_calculate_num_runners(const CRFModelConfig& model_config, - size_t batch_size, - float memory_fraction); - } // namespace dorado diff --git a/dorado/basecall/MetalCRFModel.cpp b/dorado/basecall/MetalCRFModel.cpp index a18913d2..e282c003 100644 --- a/dorado/basecall/MetalCRFModel.cpp +++ b/dorado/basecall/MetalCRFModel.cpp @@ -4,6 +4,7 @@ #include "MetalCRFModel.h" #include "CRFModelConfig.h" +#include "crf_utils.h" #include "decode/beam_search.h" #include "utils/math_utils.h" #include "utils/metal_utils.h" diff --git a/dorado/basecall/crf_utils.cpp b/dorado/basecall/crf_utils.cpp new file mode 100644 index 00000000..45c481bf --- /dev/null +++ b/dorado/basecall/crf_utils.cpp @@ -0,0 +1,76 @@ +#include "crf_utils.h" + +#include "CRFModelConfig.h" +#include "utils/memory_utils.h" +#include "utils/tensor_utils.h" + +#include +#include + +namespace dorado { +std::vector load_crf_model_weights(const std::filesystem::path &dir, + bool decomposition, + bool linear_layer_bias) { + auto tensors = std::vector{ + + "0.conv.weight.tensor", "0.conv.bias.tensor", + + "1.conv.weight.tensor", "1.conv.bias.tensor", + + "2.conv.weight.tensor", "2.conv.bias.tensor", + + "4.rnn.weight_ih_l0.tensor", "4.rnn.weight_hh_l0.tensor", + "4.rnn.bias_ih_l0.tensor", "4.rnn.bias_hh_l0.tensor", + + "5.rnn.weight_ih_l0.tensor", "5.rnn.weight_hh_l0.tensor", + "5.rnn.bias_ih_l0.tensor", "5.rnn.bias_hh_l0.tensor", + + "6.rnn.weight_ih_l0.tensor", "6.rnn.weight_hh_l0.tensor", + "6.rnn.bias_ih_l0.tensor", "6.rnn.bias_hh_l0.tensor", + + "7.rnn.weight_ih_l0.tensor", "7.rnn.weight_hh_l0.tensor", + "7.rnn.bias_ih_l0.tensor", "7.rnn.bias_hh_l0.tensor", + + "8.rnn.weight_ih_l0.tensor", "8.rnn.weight_hh_l0.tensor", + "8.rnn.bias_ih_l0.tensor", "8.rnn.bias_hh_l0.tensor", + + "9.linear.weight.tensor"}; + + if (linear_layer_bias) { + tensors.push_back("9.linear.bias.tensor"); + } + + if (decomposition) { + tensors.push_back("10.linear.weight.tensor"); + } + + return utils::load_tensors(dir, tensors); +} + +size_t auto_calculate_num_runners(const CRFModelConfig &model_config, + size_t batch_size, + float memory_fraction) { + auto model_name = std::filesystem::canonical(model_config.model_path).filename().string(); + + // very hand-wavy determination + // these numbers were determined empirically by running 1, 2, 4 and 8 runners for each model + auto required_ram_per_runner_GB = 0.f; + if (model_name.find("_fast@v") != std::string::npos) { + required_ram_per_runner_GB = 1.5; + } else if (model_name.find("_hac@v") != std::string::npos) { + required_ram_per_runner_GB = 4.5; + } else if (model_name.find("_sup@v") != std::string::npos) { + required_ram_per_runner_GB = 12.5; + } else { + return 1; + } + + // numbers were determined with a batch_size of 128, assume this just scales + required_ram_per_runner_GB *= batch_size / 128.f; + + auto free_ram_GB = utils::available_host_memory_GB() * memory_fraction; + auto num_runners = static_cast(free_ram_GB / required_ram_per_runner_GB); + return std::clamp(num_runners, size_t(1), std::size_t(std::thread::hardware_concurrency())); +} + +} // namespace dorado diff --git a/dorado/basecall/crf_utils.h b/dorado/basecall/crf_utils.h new file mode 100644 index 00000000..1652f2c9 --- /dev/null +++ b/dorado/basecall/crf_utils.h @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include +#include + +namespace dorado { + +struct CRFModelConfig; + +std::vector load_crf_model_weights(const std::filesystem::path& dir, + bool decomposition, + bool bias); + +size_t auto_calculate_num_runners(const CRFModelConfig& model_config, + size_t batch_size, + float memory_fraction); + +} // namespace dorado diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index f743195e..75ddcded 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -2,6 +2,7 @@ #include "basecall/CRFModel.h" #include "basecall/ModelRunner.h" +#include "basecall/crf_utils.h" #include "basecall/decode/CPUDecoder.h" #include "modbase/ModBaseRunner.h" From f6a459257a0fb76bd0e67096c1d4ab73b444f353 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 13:56:38 +0000 Subject: [PATCH 04/23] Move basecall functionality into `basecall` namespace --- dorado/basecall/CRFModel.cpp | 22 ++++++++++--------- dorado/basecall/CRFModel.h | 4 ++-- dorado/basecall/CRFModelConfig.cpp | 11 ++++++---- dorado/basecall/CRFModelConfig.h | 4 ++-- dorado/basecall/CudaCRFModel.cpp | 6 ++--- dorado/basecall/CudaCRFModel.h | 4 ++-- dorado/basecall/MetalCRFModel.cpp | 4 ++-- dorado/basecall/MetalCRFModel.h | 4 ++-- dorado/basecall/ModelRunner.h | 4 ++-- dorado/basecall/crf_utils.cpp | 4 ++-- dorado/basecall/crf_utils.h | 4 ++-- dorado/cli/basecaller.cpp | 2 +- dorado/cli/duplex.cpp | 12 +++++----- dorado/nn/Runners.cpp | 27 ++++++++++++----------- dorado/nn/Runners.h | 10 +++++---- dorado/read_pipeline/BasecallerNode.cpp | 4 ++-- dorado/read_pipeline/BasecallerNode.h | 4 ++-- dorado/read_pipeline/Pipelines.cpp | 16 +++++++------- dorado/read_pipeline/Pipelines.h | 11 ++++++---- dorado/read_pipeline/ScalerNode.cpp | 3 +++ dorado/read_pipeline/ScalerNode.h | 8 +++---- tests/CRFModelConfigTest.cpp | 2 +- tests/NodeSmokeTest.cpp | 29 +++++++++++++------------ 23 files changed, 107 insertions(+), 92 deletions(-) diff --git a/dorado/basecall/CRFModel.cpp b/dorado/basecall/CRFModel.cpp index 14f6727e..c0279e84 100644 --- a/dorado/basecall/CRFModel.cpp +++ b/dorado/basecall/CRFModel.cpp @@ -44,12 +44,12 @@ using Slice = torch::indexing::Slice; #if USE_KOI -KoiActivation get_koi_activation(dorado::Activation act) { - if (act == dorado::Activation::SWISH) { +static KoiActivation get_koi_activation(dorado::basecall::Activation act) { + if (act == dorado::basecall::Activation::SWISH) { return KOI_SWISH; - } else if (act == dorado::Activation::SWISH_CLAMP) { + } else if (act == dorado::basecall::Activation::SWISH_CLAMP) { return KOI_SWISH_CLAMP; - } else if (act == dorado::Activation::TANH) { + } else if (act == dorado::basecall::Activation::TANH) { return KOI_TANH; } else { throw std::logic_error("Unrecognised activation function id."); @@ -101,14 +101,15 @@ static bool koi_can_use_quantised_lstm() { return (prop->major > 6) || (prop->major == 6 && prop->minor != 2); } -static TensorLayout get_koi_lstm_input_layout(int layer_size, dorado::Activation activation) { +static TensorLayout get_koi_lstm_input_layout(int layer_size, + dorado::basecall::Activation activation) { TensorLayout layout = TensorLayout::CUBLAS_TN2C; if (koi_can_use_quantised_lstm() && (layer_size == 96 || layer_size == 128)) { layout = TensorLayout::NTC; } else if (koi_can_use_cutlass() && layer_size <= 1024 && layer_size > 128 && (layer_size % 128) == 0) { - layout = (activation == dorado::Activation::TANH) ? TensorLayout::CUTLASS_TNC_I8 - : TensorLayout::CUTLASS_TNC_F16; + layout = (activation == dorado::basecall::Activation::TANH) ? TensorLayout::CUTLASS_TNC_I8 + : TensorLayout::CUTLASS_TNC_F16; } // Apply override (Cutlass override can only be applied if conditions are met) @@ -280,7 +281,8 @@ ModuleHolder populate_model(Model &&model, const at::TensorOptions &options, bool decomposition, bool linear_layer_bias) { - auto state_dict = dorado::load_crf_model_weights(path, decomposition, linear_layer_bias); + auto state_dict = + dorado::basecall::load_crf_model_weights(path, decomposition, linear_layer_bias); model->load_state_dict(state_dict); model->to(options.dtype_opt().value().toScalarType()); model->to(options.device_opt().value()); @@ -292,7 +294,7 @@ ModuleHolder populate_model(Model &&model, } } // namespace -namespace dorado { +namespace dorado::basecall { namespace nn { @@ -913,4 +915,4 @@ ModuleHolder load_crf_model(const CRFModelConfig &model_config, model_config.out_features.has_value(), model_config.bias); } -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/CRFModel.h b/dorado/basecall/CRFModel.h index 6cc9cb58..e7045a32 100644 --- a/dorado/basecall/CRFModel.h +++ b/dorado/basecall/CRFModel.h @@ -2,11 +2,11 @@ #include -namespace dorado { +namespace dorado::basecall { struct CRFModelConfig; torch::nn::ModuleHolder load_crf_model(const CRFModelConfig& model_config, const at::TensorOptions& options); -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/CRFModelConfig.cpp b/dorado/basecall/CRFModelConfig.cpp index 35fc389e..ea787e42 100644 --- a/dorado/basecall/CRFModelConfig.cpp +++ b/dorado/basecall/CRFModelConfig.cpp @@ -12,6 +12,8 @@ #include #include +namespace { + enum SublayerType { CLAMP, CONVOLUTION, LINEAR, LINEAR_CRF_ENCODER, LSTM, PERMUTE, UNRECOGNISED }; static const std::unordered_map sublayer_map = { {"clamp", SublayerType::CLAMP}, {"convolution", SublayerType::CONVOLUTION}, @@ -29,7 +31,8 @@ SublayerType sublayer_type(const toml::value &segment) { return mapping_iter->second; } -namespace dorado { +} // namespace +namespace dorado::basecall { // Parse the config to determine if there are any clamp layers bool has_clamp(const std::vector &sublayers) { @@ -157,7 +160,7 @@ SampleType get_model_type(const std::string &model_name) { std::string SignalNormalisationParams::to_string() const { std::string str = "SignalNormalisationParams {"; - str += " strategy:" + dorado::to_string(strategy); + str += " strategy:" + dorado::basecall::to_string(strategy); if (strategy == ScalingStrategy::QUANTILE) { str += quantile.to_string(); } else if (strategy == ScalingStrategy::PA && standarisation.standardise) { @@ -173,7 +176,7 @@ std::string ConvParams::to_string() const { str += " size:" + std::to_string(size); str += " winlen:" + std::to_string(winlen); str += " stride:" + std::to_string(stride); - str += " activation:" + dorado::to_string(activation); + str += " activation:" + dorado::basecall::to_string(activation); str += "}"; return str; }; @@ -362,4 +365,4 @@ ScalingStrategy scaling_strategy_from_string(const std::string &strategy) { throw std::runtime_error("Unknown scaling strategy: `" + strategy + "`"); } -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/CRFModelConfig.h b/dorado/basecall/CRFModelConfig.h index e48d0e79..18f92c18 100644 --- a/dorado/basecall/CRFModelConfig.h +++ b/dorado/basecall/CRFModelConfig.h @@ -7,7 +7,7 @@ #include #include -namespace dorado { +namespace dorado::basecall { enum class Activation { SWISH, SWISH_CLAMP, TANH }; std::string to_string(const Activation& activation); @@ -114,4 +114,4 @@ int32_t get_model_mean_qscore_start_pos(const CRFModelConfig& model_config); bool is_rna_model(const CRFModelConfig& model_config); -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/CudaCRFModel.cpp b/dorado/basecall/CudaCRFModel.cpp index 806cf2b6..5881d30a 100644 --- a/dorado/basecall/CudaCRFModel.cpp +++ b/dorado/basecall/CudaCRFModel.cpp @@ -19,7 +19,7 @@ using namespace std::chrono_literals; -namespace dorado { +namespace dorado::basecall { class CudaCaller { public: @@ -95,7 +95,7 @@ class CudaCaller { return 64; } - int determine_batch_size(const dorado::CRFModelConfig &model_config, + int determine_batch_size(const CRFModelConfig &model_config, int chunk_size_in, float memory_limit_fraction, bool run_benchmark) { @@ -445,4 +445,4 @@ stats::NamedStats CudaModelRunner::sample_stats() const { return stats; } -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/CudaCRFModel.h b/dorado/basecall/CudaCRFModel.h index 77720cdf..902f22bd 100644 --- a/dorado/basecall/CudaCRFModel.h +++ b/dorado/basecall/CudaCRFModel.h @@ -10,7 +10,7 @@ #include #include -namespace dorado { +namespace dorado::basecall { struct CRFModelConfig; class CudaCaller; @@ -46,4 +46,4 @@ class CudaModelRunner final : public ModelRunnerBase { std::atomic m_num_batches_called = 0; }; -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/MetalCRFModel.cpp b/dorado/basecall/MetalCRFModel.cpp index e282c003..502e14bd 100644 --- a/dorado/basecall/MetalCRFModel.cpp +++ b/dorado/basecall/MetalCRFModel.cpp @@ -65,7 +65,7 @@ bool finishCommandBuffer(const char *label, MTL::CommandBuffer *cb, int try_coun } // namespace -namespace dorado { +namespace dorado::basecall { namespace nn { @@ -1098,4 +1098,4 @@ stats::NamedStats MetalModelRunner::sample_stats() const { return stats; } -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/MetalCRFModel.h b/dorado/basecall/MetalCRFModel.h index 84b5dbd9..e663d855 100644 --- a/dorado/basecall/MetalCRFModel.h +++ b/dorado/basecall/MetalCRFModel.h @@ -10,7 +10,7 @@ #include #include -namespace dorado { +namespace dorado::basecall { struct CRFModelConfig; class MetalCaller; @@ -41,4 +41,4 @@ class MetalModelRunner final : public ModelRunnerBase { std::atomic m_num_batches_called = 0; }; -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/ModelRunner.h b/dorado/basecall/ModelRunner.h index 625398fc..6b43b0e2 100644 --- a/dorado/basecall/ModelRunner.h +++ b/dorado/basecall/ModelRunner.h @@ -10,7 +10,7 @@ #include #include -namespace dorado { +namespace dorado::basecall { class ModelRunnerBase { public: @@ -109,4 +109,4 @@ stats::NamedStats ModelRunner::sample_stats() const { return stats; } -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/crf_utils.cpp b/dorado/basecall/crf_utils.cpp index 45c481bf..47ed3fd9 100644 --- a/dorado/basecall/crf_utils.cpp +++ b/dorado/basecall/crf_utils.cpp @@ -7,7 +7,7 @@ #include #include -namespace dorado { +namespace dorado::basecall { std::vector load_crf_model_weights(const std::filesystem::path &dir, bool decomposition, bool linear_layer_bias) { @@ -73,4 +73,4 @@ size_t auto_calculate_num_runners(const CRFModelConfig &model_config, return std::clamp(num_runners, size_t(1), std::size_t(std::thread::hardware_concurrency())); } -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/basecall/crf_utils.h b/dorado/basecall/crf_utils.h index 1652f2c9..57881699 100644 --- a/dorado/basecall/crf_utils.h +++ b/dorado/basecall/crf_utils.h @@ -5,7 +5,7 @@ #include #include -namespace dorado { +namespace dorado::basecall { struct CRFModelConfig; @@ -17,4 +17,4 @@ size_t auto_calculate_num_runners(const CRFModelConfig& model_config, size_t batch_size, float memory_fraction); -} // namespace dorado +} // namespace dorado::basecall diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp index e4aed1d9..5a57a01a 100644 --- a/dorado/cli/basecaller.cpp +++ b/dorado/cli/basecaller.cpp @@ -86,7 +86,7 @@ void setup(std::vector args, argparse::ArgumentParser& resume_parser, bool estimate_poly_a, const ModelSelection& model_selection) { - const auto model_config = load_crf_model_config(model_path); + const auto model_config = basecall::load_crf_model_config(model_path); const std::string model_name = models::extract_model_name_from_path(model_path); const std::string modbase_model_names = models::extract_model_names_from_paths(remora_models); diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index edcd6ba5..123a7d43 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -48,10 +48,10 @@ namespace { struct DuplexModels { std::filesystem::path model_path; std::string model_name; - CRFModelConfig model_config; + basecall::CRFModelConfig model_config; std::filesystem::path stereo_model; - CRFModelConfig stereo_model_config; + basecall::CRFModelConfig stereo_model_config; std::string stereo_model_name; std::vector mods_model_paths; @@ -122,7 +122,7 @@ DuplexModels load_models(const std::string& model_arg, } if (!skip_model_compatibility_check) { - const auto model_config = load_crf_model_config(model_path); + const auto model_config = basecall::load_crf_model_config(model_path); const auto model_name = model_path.filename().string(); check_sampling_rates_compatible(model_name, reads, model_config.sample_rate, recursive_file_loading); @@ -153,10 +153,10 @@ DuplexModels load_models(const std::string& model_arg, } const auto model_name = model_finder.get_simplex_model_name(); - const auto model_config = load_crf_model_config(model_path); + const auto model_config = basecall::load_crf_model_config(model_path); const auto stereo_model_name = stereo_model_path.filename().string(); - const auto stereo_model_config = load_crf_model_config(stereo_model_path); + const auto stereo_model_config = basecall::load_crf_model_config(stereo_model_path); return DuplexModels{model_path, model_name, model_config, stereo_model_path, @@ -488,7 +488,7 @@ int duplex(int argc, char* argv[]) { create_basecall_runners(models.model_config, device, num_runners, 0, batch_size, chunk_size, 0.9f, true); - std::vector stereo_runners; + std::vector stereo_runners; // The fraction argument for GPU memory allocates the fraction of the // _remaining_ memory to the caller. So, we allocate all of the available // memory after simplex caller has been instantiated to the duplex caller. diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index 75ddcded..2366453b 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -22,8 +22,8 @@ namespace dorado { -std::pair, size_t> create_basecall_runners( - const dorado::CRFModelConfig& model_config, +std::pair, size_t> create_basecall_runners( + const dorado::basecall::CRFModelConfig& model_config, const std::string& device, size_t num_gpu_runners, size_t num_cpu_runners, @@ -35,7 +35,7 @@ std::pair, size_t> create_basecall_runners( (void)guard_gpus; #endif - std::vector runners; + std::vector runners; // Default is 1 device. CUDA path may alter this. size_t num_devices = 1; @@ -49,22 +49,23 @@ std::pair, size_t> create_basecall_runners( batch_size = 128; } if (num_cpu_runners == 0) { - num_cpu_runners = auto_calculate_num_runners(model_config, batch_size, memory_fraction); + num_cpu_runners = + basecall::auto_calculate_num_runners(model_config, batch_size, memory_fraction); } spdlog::debug("- CPU calling: set batch size to {}, num_cpu_runners to {}", batch_size, num_cpu_runners); for (size_t i = 0; i < num_cpu_runners; i++) { - runners.push_back(std::make_shared>( + runners.push_back(std::make_shared>( model_config, device, int(chunk_size), int(batch_size))); } } #if DORADO_GPU_BUILD #ifdef __APPLE__ else if (device == "metal") { - auto caller = dorado::create_metal_caller(model_config, chunk_size, batch_size); + auto caller = basecall::create_metal_caller(model_config, chunk_size, batch_size); for (size_t i = 0; i < num_gpu_runners; i++) { - runners.push_back(std::make_shared(caller)); + runners.push_back(std::make_shared(caller)); } if (batch_size == 0) { spdlog::info(" - set batch size to {}", runners.back()->batch_size()); @@ -85,13 +86,13 @@ std::pair, size_t> create_basecall_runners( } cxxpool::thread_pool pool{num_devices}; - std::vector> callers; - std::vector>> futures; + std::vector> callers; + std::vector>> futures; for (auto device_string : devices) { - futures.push_back(pool.push(dorado::create_cuda_caller, model_config, int(chunk_size), - int(batch_size), device_string, memory_fraction, - guard_gpus)); + futures.push_back(pool.push(dorado::basecall::create_cuda_caller, model_config, + int(chunk_size), int(batch_size), device_string, + memory_fraction, guard_gpus)); } for (auto& caller : futures) { @@ -100,7 +101,7 @@ std::pair, size_t> create_basecall_runners( for (size_t j = 0; j < num_devices; j++) { for (size_t i = 0; i < num_gpu_runners; i++) { - runners.push_back(std::make_shared(callers[j])); + runners.push_back(std::make_shared(callers[j])); } if (batch_size == 0) { spdlog::info(" - set batch size for {} to {}", devices[j], diff --git a/dorado/nn/Runners.h b/dorado/nn/Runners.h index 9276fdf2..dca26e64 100644 --- a/dorado/nn/Runners.h +++ b/dorado/nn/Runners.h @@ -8,14 +8,16 @@ namespace dorado { +namespace basecall { struct CRFModelConfig; class ModelRunnerBase; -class ModBaseRunner; - using Runner = std::shared_ptr; +} // namespace basecall + +class ModBaseRunner; -std::pair, size_t> create_basecall_runners( - const dorado::CRFModelConfig& model_config, +std::pair, size_t> create_basecall_runners( + const dorado::basecall::CRFModelConfig& model_config, const std::string& device, size_t num_gpu_runners, size_t num_cpu_runners, diff --git a/dorado/read_pipeline/BasecallerNode.cpp b/dorado/read_pipeline/BasecallerNode.cpp index 7048029c..6d8f72fa 100644 --- a/dorado/read_pipeline/BasecallerNode.cpp +++ b/dorado/read_pipeline/BasecallerNode.cpp @@ -288,7 +288,7 @@ void BasecallerNode::basecall_worker_thread(int worker_id) { namespace { // Calculates the input queue size. -size_t CalcMaxChunksIn(const std::vector &model_runners) { +size_t CalcMaxChunksIn(const std::vector &model_runners) { // Allow 2 batches per model runner on the chunks_in queue size_t max_chunks_in = 0; // Allows optimal batch size to be used for every GPU @@ -300,7 +300,7 @@ size_t CalcMaxChunksIn(const std::vector &model_runners) { } // namespace -BasecallerNode::BasecallerNode(std::vector model_runners, +BasecallerNode::BasecallerNode(std::vector model_runners, size_t overlap, int batch_timeout_ms, std::string model_name, diff --git a/dorado/read_pipeline/BasecallerNode.h b/dorado/read_pipeline/BasecallerNode.h index 9f0f9862..f99dc5f1 100644 --- a/dorado/read_pipeline/BasecallerNode.h +++ b/dorado/read_pipeline/BasecallerNode.h @@ -20,7 +20,7 @@ class BasecallerNode : public MessageSink { public: // Chunk size and overlap are in raw samples - BasecallerNode(std::vector model_runners, + BasecallerNode(std::vector model_runners, size_t overlap, int batch_timeout_ms, std::string model_name, @@ -46,7 +46,7 @@ class BasecallerNode : public MessageSink { void working_reads_manager(); // Vector of model runners (each with their own GPU access etc) - std::vector m_model_runners; + std::vector m_model_runners; // Chunk length size_t m_chunk_size; // Minimum overlap between two adjacent chunks in a read. Overlap is used to reduce edge effects and improve accuracy. diff --git a/dorado/read_pipeline/Pipelines.cpp b/dorado/read_pipeline/Pipelines.cpp index 435fa4b2..57d14b27 100644 --- a/dorado/read_pipeline/Pipelines.cpp +++ b/dorado/read_pipeline/Pipelines.cpp @@ -17,7 +17,7 @@ namespace dorado::pipelines { void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, + std::vector&& runners, std::vector>&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, @@ -76,7 +76,7 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, // For DNA, read splitting happens after basecall. if (enable_read_splitter && !is_rna) { splitter::DuplexSplitSettings splitter_settings(model_config.signal_norm_params.strategy == - ScalingStrategy::PA); + basecall::ScalingStrategy::PA); splitter_settings.simplex_mode = true; auto dna_splitter = std::make_unique(splitter_settings); auto dna_splitter_node = pipeline_desc.add_node({}, std::move(dna_splitter), @@ -107,8 +107,8 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, void create_stereo_duplex_pipeline( PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, + std::vector&& runners, + std::vector&& stereo_runners, std::vector>&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, @@ -162,7 +162,7 @@ void create_stereo_duplex_pipeline( // as a passthrough, meaning it won't perform any splitting operations and // will just pass data through. splitter::DuplexSplitSettings splitter_settings(model_config.signal_norm_params.strategy == - ScalingStrategy::PA); + basecall::ScalingStrategy::PA); auto duplex_splitter = std::make_unique(splitter_settings); auto splitter_node = pipeline_desc.add_node( {pairing_node}, std::move(duplex_splitter), splitter_node_threads, 1000); @@ -174,9 +174,9 @@ void create_stereo_duplex_pipeline( {splitter_node}, std::move(runners), adjusted_simplex_overlap, kSimplexBatchTimeoutMS, model_name, 1000, "BasecallerNode", mean_qscore_start_pos); - auto scaler_node = - pipeline_desc.add_node({basecaller_node}, model_config.signal_norm_params, - SampleType::DNA, scaler_node_threads, 1000); + auto scaler_node = pipeline_desc.add_node( + {basecaller_node}, model_config.signal_norm_params, basecall::SampleType::DNA, + scaler_node_threads, 1000); // if we've been provided a source node, connect it to the start of our pipeline if (source_node_handle != PipelineDescriptor::InvalidNodeHandle) { diff --git a/dorado/read_pipeline/Pipelines.h b/dorado/read_pipeline/Pipelines.h index f4a27211..e419dcc9 100644 --- a/dorado/read_pipeline/Pipelines.h +++ b/dorado/read_pipeline/Pipelines.h @@ -12,9 +12,12 @@ namespace dorado { class ModBaseRunner; -class ModelRunnerBase; +namespace basecall { +class ModelRunnerBase; using Runner = std::shared_ptr; +} // namespace basecall + using PairingParameters = std::variant>; namespace pipelines { @@ -23,7 +26,7 @@ namespace pipelines { /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, + std::vector&& runners, std::vector>&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, @@ -39,8 +42,8 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline void create_stereo_duplex_pipeline( PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, + std::vector&& runners, + std::vector&& stereo_runners, std::vector>&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, diff --git a/dorado/read_pipeline/ScalerNode.cpp b/dorado/read_pipeline/ScalerNode.cpp index 2be915e6..9e1a706c 100644 --- a/dorado/read_pipeline/ScalerNode.cpp +++ b/dorado/read_pipeline/ScalerNode.cpp @@ -20,6 +20,9 @@ using namespace std::chrono_literals; using Slice = at::indexing::Slice; namespace dorado { +using SampleType = dorado::basecall::SampleType; +using ScalingStrategy = dorado::basecall::ScalingStrategy; +using SignalNormalisationParams = dorado::basecall::SignalNormalisationParams; std::pair ScalerNode::normalisation(const at::Tensor& x) { // Calculate shift and scale factors for normalisation. diff --git a/dorado/read_pipeline/ScalerNode.h b/dorado/read_pipeline/ScalerNode.h index 4bd521a0..e4cea5ef 100644 --- a/dorado/read_pipeline/ScalerNode.h +++ b/dorado/read_pipeline/ScalerNode.h @@ -15,8 +15,8 @@ namespace dorado { class ScalerNode : public MessageSink { public: - ScalerNode(const SignalNormalisationParams& config, - SampleType model_type, + ScalerNode(const basecall::SignalNormalisationParams& config, + basecall::SampleType model_type, int num_worker_threads, size_t max_reads); ~ScalerNode() { terminate_impl(); } @@ -32,8 +32,8 @@ class ScalerNode : public MessageSink { std::vector> m_worker_threads; std::atomic m_num_worker_threads; - SignalNormalisationParams m_scaling_params; - const SampleType m_model_type; + basecall::SignalNormalisationParams m_scaling_params; + const basecall::SampleType m_model_type; std::pair med_mad(const at::Tensor& x); std::pair normalisation(const at::Tensor& x); diff --git a/tests/CRFModelConfigTest.cpp b/tests/CRFModelConfigTest.cpp index 6a47224d..1333f226 100644 --- a/tests/CRFModelConfigTest.cpp +++ b/tests/CRFModelConfigTest.cpp @@ -12,7 +12,7 @@ #define CUT_TAG "[CRFModelConfig]" -using namespace dorado; +using namespace dorado::basecall; namespace fs = std::filesystem; TEST_CASE(CUT_TAG ": test dna_r9.4.1 hac@v3.3 model load", CUT_TAG) { diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 881ff9d1..39e3f64d 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -165,8 +165,9 @@ TempDir download_model(const std::string& model) { DEFINE_TEST(NodeSmokeTestRead, "ScalerNode") { auto pipeline_restart = GENERATE(false, true); - auto model_type = GENERATE(dorado::SampleType::DNA, dorado::SampleType::RNA002, - dorado::SampleType::RNA004); + auto model_type = + GENERATE(dorado::basecall::SampleType::DNA, dorado::basecall::SampleType::RNA002, + dorado::basecall::SampleType::RNA004); CAPTURE(pipeline_restart); CAPTURE(model_type); @@ -177,8 +178,8 @@ DEFINE_TEST(NodeSmokeTestRead, "ScalerNode") { read->read_common.raw_data = read->read_common.raw_data.to(torch::kI16); }); - dorado::SignalNormalisationParams config; - config.strategy = dorado::ScalingStrategy::QUANTILE; + dorado::basecall::SignalNormalisationParams config; + config.strategy = dorado::basecall::ScalingStrategy::QUANTILE; config.quantile.quantile_a = 0.2f; config.quantile.quantile_b = 0.9f; config.quantile.shift_multiplier = 0.51f; @@ -202,20 +203,20 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { const auto& default_params = dorado::utils::default_parameters; const auto model_dir = download_model(model_name); const auto model_path = (model_dir.m_path / model_name).string(); - auto model_config = dorado::load_crf_model_config(model_path); + auto model_config = dorado::basecall::load_crf_model_config(model_path); // Use a fixed batch size otherwise we slow down CI autobatchsizing. std::size_t batch_size = 128; // Create runners - std::vector runners; + std::vector runners; if (gpu) { #if DORADO_GPU_BUILD #ifdef __APPLE__ - auto caller = - dorado::create_metal_caller(model_config, default_params.chunksize, batch_size); + auto caller = dorado::basecall::create_metal_caller(model_config, default_params.chunksize, + batch_size); for (int i = 0; i < default_params.num_runners; i++) { - runners.push_back(std::make_shared(caller)); + runners.push_back(std::make_shared(caller)); } #else // __APPLE__ auto devices = dorado::utils::parse_cuda_device_string("cuda:all"); @@ -223,10 +224,10 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { SKIP("No CUDA devices found"); } for (const auto& device : devices) { - auto caller = dorado::create_cuda_caller(model_config, default_params.chunksize, - int(batch_size), device, 1.f, true); + auto caller = dorado::basecall::create_cuda_caller( + model_config, default_params.chunksize, int(batch_size), device, 1.f, true); for (int i = 0; i < default_params.num_runners; i++) { - runners.push_back(std::make_shared(caller)); + runners.push_back(std::make_shared(caller)); } } #endif // __APPLE__ @@ -238,7 +239,7 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { set_num_reads(5); set_expected_messages(5); batch_size = 8; - runners.push_back(std::make_shared>( + runners.push_back(std::make_shared>( model_config, "cpu", default_params.chunksize, int(batch_size))); } @@ -272,7 +273,7 @@ DEFINE_TEST(NodeSmokeTestRead, "ModBaseCallerNode") { const char model_name[] = "dna_r10.4.1_e8.2_400bps_fast@v4.2.0"; const auto model_dir = download_model(model_name); const std::size_t model_stride = - dorado::load_crf_model_config(model_dir.m_path / model_name).stride; + dorado::basecall::load_crf_model_config(model_dir.m_path / model_name).stride; // Create runners std::vector> remora_runners; From e15ebf731853076e00d825fef2337c49a2b6b259 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 14:17:03 +0000 Subject: [PATCH 05/23] Move modbase functionality into `modbase` namespace --- dorado/modbase/ModBaseContext.cpp | 14 +++++----- dorado/modbase/ModBaseContext.h | 7 ++--- dorado/modbase/ModBaseModel.cpp | 4 +-- dorado/modbase/ModBaseModel.h | 4 +-- dorado/modbase/ModBaseModelConfig.cpp | 4 +-- dorado/modbase/ModBaseModelConfig.h | 4 +-- dorado/modbase/ModBaseRunner.cpp | 4 +-- dorado/modbase/ModBaseRunner.h | 4 +-- dorado/modbase/ModbaseEncoder.cpp | 4 +-- dorado/modbase/ModbaseEncoder.h | 4 +-- dorado/modbase/ModbaseScaler.cpp | 4 +-- dorado/modbase/ModbaseScaler.h | 4 +-- dorado/modbase/MotifMatcher.cpp | 4 +-- dorado/modbase/MotifMatcher.h | 4 +-- dorado/nn/Runners.cpp | 10 +++---- dorado/nn/Runners.h | 4 ++- dorado/read_pipeline/ModBaseCallerNode.cpp | 23 ++++++++-------- dorado/read_pipeline/ModBaseCallerNode.h | 7 ++--- dorado/read_pipeline/Pipelines.cpp | 25 ++++++++--------- dorado/read_pipeline/Pipelines.h | 31 ++++++++++++---------- dorado/read_pipeline/ReadPipeline.cpp | 2 +- tests/ModBaseEncoderTest.cpp | 2 +- tests/MotifMatcherTest.cpp | 4 +-- tests/NodeSmokeTest.cpp | 8 +++--- 24 files changed, 96 insertions(+), 89 deletions(-) diff --git a/dorado/modbase/ModBaseContext.cpp b/dorado/modbase/ModBaseContext.cpp index b7845be5..c38a579b 100644 --- a/dorado/modbase/ModBaseContext.cpp +++ b/dorado/modbase/ModBaseContext.cpp @@ -5,14 +5,16 @@ #include -namespace dorado::utils { +namespace dorado::modbase { ModBaseContext::ModBaseContext() {} ModBaseContext::~ModBaseContext() {} -const std::string& ModBaseContext::motif(char base) const { return m_motifs[base_to_int(base)]; } +const std::string& ModBaseContext::motif(char base) const { + return m_motifs[utils::base_to_int(base)]; +} -size_t ModBaseContext::motif_offset(char base) const { return m_offsets[base_to_int(base)]; } +size_t ModBaseContext::motif_offset(char base) const { return m_offsets[utils::base_to_int(base)]; } void ModBaseContext::set_context(std::string motif, size_t offset) { if (motif.size() < 2) { @@ -20,7 +22,7 @@ void ModBaseContext::set_context(std::string motif, size_t offset) { return; } char base = motif.at(offset); - auto index = base_to_int(base); + auto index = utils::base_to_int(base); m_motif_matchers[index] = std::make_unique(motif, offset); m_motifs[index] = std::move(motif); m_offsets[index] = offset; @@ -99,7 +101,7 @@ void ModBaseContext::update_mask(std::vector& mask, // A cardinal base. current_cardinal = modbase_alphabet[channel_idx][0]; } else { - if (!m_motifs[base_to_int(current_cardinal)].empty()) { + if (!m_motifs[utils::base_to_int(current_cardinal)].empty()) { // This cardinal base has a context associated with modifications, so the mask should // not be updated, regardless of the threshold. continue; @@ -115,4 +117,4 @@ void ModBaseContext::update_mask(std::vector& mask, } } -} // namespace dorado::utils \ No newline at end of file +} // namespace dorado::modbase \ No newline at end of file diff --git a/dorado/modbase/ModBaseContext.h b/dorado/modbase/ModBaseContext.h index c806091f..75e06166 100644 --- a/dorado/modbase/ModBaseContext.h +++ b/dorado/modbase/ModBaseContext.h @@ -24,12 +24,9 @@ * be "CXT:XG:_:_". */ -namespace dorado { +namespace dorado::modbase { struct ModBaseModelConfig; class MotifMatcher; -} // namespace dorado - -namespace dorado::utils { class ModBaseContext { public: @@ -97,4 +94,4 @@ class ModBaseContext { std::array, 4> m_motif_matchers; }; -} // namespace dorado::utils +} // namespace dorado::modbase diff --git a/dorado/modbase/ModBaseModel.cpp b/dorado/modbase/ModBaseModel.cpp index de1a03be..1ccae387 100644 --- a/dorado/modbase/ModBaseModel.cpp +++ b/dorado/modbase/ModBaseModel.cpp @@ -28,7 +28,7 @@ ModuleHolder populate_model(Model&& model, } } // namespace -namespace dorado { +namespace dorado::modbase { namespace nn { @@ -257,4 +257,4 @@ ModuleHolder load_modbase_model(const std::filesystem::path& model_pa throw std::runtime_error("Unknown model type in config file."); } -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModBaseModel.h b/dorado/modbase/ModBaseModel.h index b1ce4f80..a8c6111f 100644 --- a/dorado/modbase/ModBaseModel.h +++ b/dorado/modbase/ModBaseModel.h @@ -4,10 +4,10 @@ #include -namespace dorado { +namespace dorado::modbase { torch::nn::ModuleHolder load_modbase_model( const std::filesystem::path& model_path, at::TensorOptions options); -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModBaseModelConfig.cpp b/dorado/modbase/ModBaseModelConfig.cpp index 711bae41..1111c7fd 100644 --- a/dorado/modbase/ModBaseModelConfig.cpp +++ b/dorado/modbase/ModBaseModelConfig.cpp @@ -6,7 +6,7 @@ #include -namespace dorado { +namespace dorado::modbase { ModBaseModelConfig load_modbase_model_config(const std::filesystem::path& model_path) { ModBaseModelConfig config; @@ -130,4 +130,4 @@ ModBaseInfo get_modbase_info( return result; } -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModBaseModelConfig.h b/dorado/modbase/ModBaseModelConfig.h index cd7d9a22..84fabbdc 100644 --- a/dorado/modbase/ModBaseModelConfig.h +++ b/dorado/modbase/ModBaseModelConfig.h @@ -6,7 +6,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { struct ModBaseModelConfig { std::vector mod_long_names; ///< The long names of the modified bases. @@ -32,4 +32,4 @@ ModBaseModelConfig load_modbase_model_config(const std::filesystem::path & model ModBaseInfo get_modbase_info( const std::vector> & base_mod_params); -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModBaseRunner.cpp b/dorado/modbase/ModBaseRunner.cpp index b55f9bf9..4fdc8ea2 100644 --- a/dorado/modbase/ModBaseRunner.cpp +++ b/dorado/modbase/ModBaseRunner.cpp @@ -22,7 +22,7 @@ using namespace std::chrono_literals; -namespace dorado { +namespace dorado::modbase { class ModBaseCaller { public: @@ -354,4 +354,4 @@ stats::NamedStats ModBaseRunner::sample_stats() const { return stats; } -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModBaseRunner.h b/dorado/modbase/ModBaseRunner.h index 04f1842d..1de58d35 100644 --- a/dorado/modbase/ModBaseRunner.h +++ b/dorado/modbase/ModBaseRunner.h @@ -9,7 +9,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { struct ModBaseModelConfig; class ModBaseCaller; @@ -49,4 +49,4 @@ class ModBaseRunner { std::atomic m_num_batches_called = 0; }; -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModbaseEncoder.cpp b/dorado/modbase/ModbaseEncoder.cpp index 77f6943d..e98f9938 100644 --- a/dorado/modbase/ModbaseEncoder.cpp +++ b/dorado/modbase/ModbaseEncoder.cpp @@ -9,7 +9,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { ModBaseEncoder::ModBaseEncoder(size_t block_stride, size_t context_samples, @@ -275,4 +275,4 @@ std::vector ModBaseEncoder::encode_kmer(const std::vector& seq, m_kmer_len); } -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModbaseEncoder.h b/dorado/modbase/ModbaseEncoder.h index 3360aab3..6d5e7dc7 100644 --- a/dorado/modbase/ModbaseEncoder.h +++ b/dorado/modbase/ModbaseEncoder.h @@ -4,7 +4,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { class ModBaseEncoder { private: @@ -61,4 +61,4 @@ class ModBaseEncoder { Context get_context(size_t seq_pos) const; }; -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModbaseScaler.cpp b/dorado/modbase/ModbaseScaler.cpp index 7817fede..c8c3971a 100644 --- a/dorado/modbase/ModbaseScaler.cpp +++ b/dorado/modbase/ModbaseScaler.cpp @@ -8,7 +8,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { ModBaseScaler::ModBaseScaler(const std::vector& kmer_levels, size_t kmer_len, @@ -98,4 +98,4 @@ std::pair ModBaseScaler::calc_offset_scale( return {new_offset, new_scale}; } -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/ModbaseScaler.h b/dorado/modbase/ModbaseScaler.h index eb3db861..2c8e1304 100644 --- a/dorado/modbase/ModbaseScaler.h +++ b/dorado/modbase/ModbaseScaler.h @@ -7,7 +7,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { /// Calculates new scaling values for improved modified base detection class ModBaseScaler { @@ -64,4 +64,4 @@ class ModBaseScaler { ModBaseScaler(const std::vector& kmer_levels, size_t kmer_len, size_t centre_index); }; -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/MotifMatcher.cpp b/dorado/modbase/MotifMatcher.cpp index 4570d939..ec601293 100644 --- a/dorado/modbase/MotifMatcher.cpp +++ b/dorado/modbase/MotifMatcher.cpp @@ -41,7 +41,7 @@ std::string expand_motif_regex(const std::string& motif) { } // namespace -namespace dorado { +namespace dorado::modbase { MotifMatcher::MotifMatcher(const ModBaseModelConfig& model_config) : MotifMatcher(model_config.motif, model_config.motif_offset) {} @@ -68,4 +68,4 @@ std::vector MotifMatcher::get_motif_hits(std::string_view seq) const { return context_hits; } -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/modbase/MotifMatcher.h b/dorado/modbase/MotifMatcher.h index f203f3cf..0501a34b 100644 --- a/dorado/modbase/MotifMatcher.h +++ b/dorado/modbase/MotifMatcher.h @@ -4,7 +4,7 @@ #include #include -namespace dorado { +namespace dorado::modbase { struct ModBaseModelConfig; class MotifMatcher { @@ -19,4 +19,4 @@ class MotifMatcher { const size_t m_motif_offset; }; -} // namespace dorado +} // namespace dorado::modbase diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index 2366453b..2f22b4ec 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -137,7 +137,7 @@ std::pair, size_t> create_basecall_runners return {runners, num_devices}; } -std::vector> create_modbase_runners( +std::vector> create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, @@ -147,7 +147,7 @@ std::vector> create_modbase_runners( } // generate model callers before nodes or it affects the speed calculations - std::vector> remora_runners; + std::vector> remora_runners; std::vector modbase_devices; int remora_callers = 1; @@ -170,10 +170,10 @@ std::vector> create_modbase_runners( #endif // DORADO_GPU_BUILD for (const auto& device_string : modbase_devices) { for (int i = 0; i < remora_callers; ++i) { - auto caller = dorado::create_modbase_caller(remora_models, int(remora_batch_size), - device_string); + auto caller = dorado::modbase::create_modbase_caller( + remora_models, int(remora_batch_size), device_string); for (size_t j = 0; j < remora_runners_per_caller; j++) { - remora_runners.push_back(std::make_unique(caller)); + remora_runners.push_back(std::make_unique(caller)); } } }; diff --git a/dorado/nn/Runners.h b/dorado/nn/Runners.h index dca26e64..9ca6624c 100644 --- a/dorado/nn/Runners.h +++ b/dorado/nn/Runners.h @@ -14,7 +14,9 @@ class ModelRunnerBase; using Runner = std::shared_ptr; } // namespace basecall +namespace modbase { class ModBaseRunner; +} std::pair, size_t> create_basecall_runners( const dorado::basecall::CRFModelConfig& model_config, @@ -26,7 +28,7 @@ std::pair, size_t> create_basecall_runners float memory_fraction, bool guard_gpus); -std::vector> create_modbase_runners( +std::vector> create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, diff --git a/dorado/read_pipeline/ModBaseCallerNode.cpp b/dorado/read_pipeline/ModBaseCallerNode.cpp index ecdce869..e5dfd58b 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.cpp +++ b/dorado/read_pipeline/ModBaseCallerNode.cpp @@ -50,10 +50,11 @@ struct ModBaseCallerNode::WorkingRead { num_modbase_chunks_called; // Number of modbase chunks which have been scored }; -ModBaseCallerNode::ModBaseCallerNode(std::vector> model_runners, - size_t remora_threads, - size_t block_stride, - size_t max_reads) +ModBaseCallerNode::ModBaseCallerNode( + std::vector> model_runners, + size_t remora_threads, + size_t block_stride, + size_t max_reads) : MessageSink(max_reads), m_runners(std::move(model_runners)), m_num_input_workers(remora_threads), @@ -125,9 +126,9 @@ void ModBaseCallerNode::restart() { } void ModBaseCallerNode::init_modbase_info() { - std::vector> base_mod_params; + std::vector> base_mod_params; auto& runner = m_runners[0]; - utils::ModBaseContext context_handler; + modbase::ModBaseContext context_handler; for (size_t caller_id = 0; caller_id < runner->num_callers(); ++caller_id) { const auto& params = runner->caller_params(caller_id); if (!params.motif.empty()) { @@ -137,7 +138,7 @@ void ModBaseCallerNode::init_modbase_info() { m_num_states += params.base_mod_count; } - auto result = get_modbase_info(base_mod_params); + auto result = modbase::get_modbase_info(base_mod_params); m_mod_base_info = std::make_shared( std::move(result.alphabet), std::move(result.long_names), context_handler.encode()); @@ -229,8 +230,8 @@ void ModBaseCallerNode::duplex_mod_call(Message&& message) { auto context_samples = (params.context_before + params.context_after); // One-hot encodes the kmer at each signal step for input into the network - ModBaseEncoder encoder(m_block_stride, context_samples, params.bases_before, - params.bases_after); + modbase::ModBaseEncoder encoder(m_block_stride, context_samples, + params.bases_before, params.bases_after); encoder.init(sequence_ints, seq_to_sig_map); auto context_hits = runner->get_motif_hits(caller_id, new_seq); @@ -355,8 +356,8 @@ void ModBaseCallerNode::simplex_mod_call(Message&& message) { auto context_samples = (params.context_before + params.context_after); // One-hot encodes the kmer at each signal step for input into the network - ModBaseEncoder encoder(m_block_stride, context_samples, params.bases_before, - params.bases_after); + modbase::ModBaseEncoder encoder(m_block_stride, context_samples, params.bases_before, + params.bases_after); encoder.init(sequence_ints, seq_to_sig_map); auto context_hits = runner->get_motif_hits(caller_id, read->read_common.seq); diff --git a/dorado/read_pipeline/ModBaseCallerNode.h b/dorado/read_pipeline/ModBaseCallerNode.h index a97d032f..968c0523 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.h +++ b/dorado/read_pipeline/ModBaseCallerNode.h @@ -16,14 +16,15 @@ namespace dorado { +namespace modbase { class ModBaseRunner; - +} class ModBaseCallerNode : public MessageSink { struct RemoraChunk; struct WorkingRead; public: - ModBaseCallerNode(std::vector> model_runners, + ModBaseCallerNode(std::vector> model_runners, size_t remora_threads, size_t block_stride, size_t max_reads); @@ -56,7 +57,7 @@ class ModBaseCallerNode : public MessageSink { // Worker thread, processes chunk results back into the reads void output_worker_thread(); - std::vector> m_runners; + std::vector> m_runners; size_t m_num_input_workers = 0; size_t m_block_stride; size_t m_batch_size; diff --git a/dorado/read_pipeline/Pipelines.cpp b/dorado/read_pipeline/Pipelines.cpp index 57d14b27..bb19a51f 100644 --- a/dorado/read_pipeline/Pipelines.cpp +++ b/dorado/read_pipeline/Pipelines.cpp @@ -16,17 +16,18 @@ namespace dorado::pipelines { -void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector>&& modbase_runners, - size_t overlap, - uint32_t mean_qscore_start_pos, - int scaler_node_threads, - bool enable_read_splitter, - int splitter_node_threads, - int modbase_node_threads, - NodeHandle sink_node_handle, - NodeHandle source_node_handle) { +void create_simplex_pipeline( + PipelineDescriptor& pipeline_desc, + std::vector&& runners, + std::vector>&& modbase_runners, + size_t overlap, + uint32_t mean_qscore_start_pos, + int scaler_node_threads, + bool enable_read_splitter, + int splitter_node_threads, + int modbase_node_threads, + NodeHandle sink_node_handle, + NodeHandle source_node_handle) { const auto& model_config = runners.front()->config(); auto model_stride = runners.front()->model_stride(); auto adjusted_overlap = (overlap / model_stride) * model_stride; @@ -109,7 +110,7 @@ void create_stereo_duplex_pipeline( PipelineDescriptor& pipeline_desc, std::vector&& runners, std::vector&& stereo_runners, - std::vector>&& modbase_runners, + std::vector>&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, diff --git a/dorado/read_pipeline/Pipelines.h b/dorado/read_pipeline/Pipelines.h index e419dcc9..d68dea4a 100644 --- a/dorado/read_pipeline/Pipelines.h +++ b/dorado/read_pipeline/Pipelines.h @@ -11,13 +11,15 @@ namespace dorado { -class ModBaseRunner; - namespace basecall { class ModelRunnerBase; using Runner = std::shared_ptr; } // namespace basecall +namespace modbase { +class ModBaseRunner; +} + using PairingParameters = std::variant>; namespace pipelines { @@ -25,17 +27,18 @@ namespace pipelines { /// Create a simplex basecall pipeline description /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline -void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector>&& modbase_runners, - size_t overlap, - uint32_t mean_qscore_start_pos, - int scaler_node_threads, - bool enable_read_splitter, - int splitter_node_threads, - int modbase_threads, - NodeHandle sink_node_handle, - NodeHandle source_node_handle); +void create_simplex_pipeline( + PipelineDescriptor& pipeline_desc, + std::vector&& runners, + std::vector>&& modbase_runners, + size_t overlap, + uint32_t mean_qscore_start_pos, + int scaler_node_threads, + bool enable_read_splitter, + int splitter_node_threads, + int modbase_threads, + NodeHandle sink_node_handle, + NodeHandle source_node_handle); /// Create a duplex basecall pipeline description /// If source_node_handle is valid, set this to be the source of the simplex pipeline @@ -44,7 +47,7 @@ void create_stereo_duplex_pipeline( PipelineDescriptor& pipeline_desc, std::vector&& runners, std::vector&& stereo_runners, - std::vector>&& modbase_runners, + std::vector>&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, diff --git a/dorado/read_pipeline/ReadPipeline.cpp b/dorado/read_pipeline/ReadPipeline.cpp index 42b7f668..d79371a0 100644 --- a/dorado/read_pipeline/ReadPipeline.cpp +++ b/dorado/read_pipeline/ReadPipeline.cpp @@ -158,7 +158,7 @@ void ReadCommon::generate_modbase_tags(bam1_t *aln, uint8_t threshold) const { // Create a mask indicating which bases are modified. std::unordered_map base_has_context = { {'A', false}, {'C', false}, {'G', false}, {'T', false}}; - utils::ModBaseContext context_handler; + modbase::ModBaseContext context_handler; if (!mod_base_info->context.empty()) { if (!context_handler.decode(mod_base_info->context)) { throw std::runtime_error("Invalid base modification context string."); diff --git a/tests/ModBaseEncoderTest.cpp b/tests/ModBaseEncoderTest.cpp index 6ed0dd07..feace473 100644 --- a/tests/ModBaseEncoderTest.cpp +++ b/tests/ModBaseEncoderTest.cpp @@ -16,7 +16,7 @@ TEST_CASE("Encode sequence for modified basecalling", TEST_GROUP) { auto seq_to_sig_map = dorado::utils::moves_to_map(moves, BLOCK_STRIDE, moves.size() * BLOCK_STRIDE, std::nullopt); - dorado::ModBaseEncoder encoder(BLOCK_STRIDE, SLICE_BLOCKS * BLOCK_STRIDE, 1, 1); + dorado::modbase::ModBaseEncoder encoder(BLOCK_STRIDE, SLICE_BLOCKS * BLOCK_STRIDE, 1, 1); encoder.init(seq_ints, seq_to_sig_map); auto slice0 = encoder.get_context(0); // The T in the NTA 3mer. diff --git a/tests/MotifMatcherTest.cpp b/tests/MotifMatcherTest.cpp index d62207d7..ad9ed2bc 100644 --- a/tests/MotifMatcherTest.cpp +++ b/tests/MotifMatcherTest.cpp @@ -24,7 +24,7 @@ const std::string SEQ = "AACCGGTTACGTGGACTGACACTAAA"; } // namespace TEST_CASE(TEST_GROUP ": test motifs", TEST_GROUP) { - dorado::ModBaseModelConfig config; + dorado::modbase::ModBaseModelConfig config; auto [motif, motif_offset, expected_results] = GENERATE(table>({ // clang-format off @@ -42,7 +42,7 @@ TEST_CASE(TEST_GROUP ": test motifs", TEST_GROUP) { config.motif = motif; config.motif_offset = motif_offset; - dorado::MotifMatcher matcher(config); + dorado::modbase::MotifMatcher matcher(config); auto hits = matcher.get_motif_hits(SEQ); CHECK(hits == expected_results); } diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 39e3f64d..dadac3e6 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -276,7 +276,7 @@ DEFINE_TEST(NodeSmokeTestRead, "ModBaseCallerNode") { dorado::basecall::load_crf_model_config(model_dir.m_path / model_name).stride; // Create runners - std::vector> remora_runners; + std::vector> remora_runners; std::vector modbase_devices; int batch_size = default_params.remora_batchsize; if (gpu) { @@ -300,10 +300,10 @@ DEFINE_TEST(NodeSmokeTestRead, "ModBaseCallerNode") { batch_size = 8; // reduce batch size so we're not doing work on empty entries } for (const auto& device_string : modbase_devices) { - auto caller = dorado::create_modbase_caller({remora_model, remora_model_6mA}, batch_size, - device_string); + auto caller = dorado::modbase::create_modbase_caller({remora_model, remora_model_6mA}, + batch_size, device_string); for (int i = 0; i < default_params.mod_base_runners_per_caller; i++) { - remora_runners.push_back(std::make_unique(caller)); + remora_runners.push_back(std::make_unique(caller)); } } From 2df69adcfef02b268e90343097d26615ceab3baa Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 14:57:59 +0000 Subject: [PATCH 06/23] Tidy use of modbase runners --- dorado/modbase/ModBaseRunner.h | 2 + dorado/nn/Runners.cpp | 4 +- dorado/nn/Runners.h | 5 ++- dorado/read_pipeline/ModBaseCallerNode.cpp | 9 ++-- dorado/read_pipeline/ModBaseCallerNode.h | 8 ++-- dorado/read_pipeline/Pipelines.cpp | 48 ++++++++++---------- dorado/read_pipeline/Pipelines.h | 51 +++++++++++----------- tests/NodeSmokeTest.cpp | 2 +- 8 files changed, 65 insertions(+), 64 deletions(-) diff --git a/dorado/modbase/ModBaseRunner.h b/dorado/modbase/ModBaseRunner.h index 1de58d35..a1e0fa3d 100644 --- a/dorado/modbase/ModBaseRunner.h +++ b/dorado/modbase/ModBaseRunner.h @@ -49,4 +49,6 @@ class ModBaseRunner { std::atomic m_num_batches_called = 0; }; +using Runner = std::unique_ptr; + } // namespace dorado::modbase diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index 2f22b4ec..c3e5755c 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -137,7 +137,7 @@ std::pair, size_t> create_basecall_runners return {runners, num_devices}; } -std::vector> create_modbase_runners( +std::vector create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, @@ -147,7 +147,7 @@ std::vector> create_modbase_runn } // generate model callers before nodes or it affects the speed calculations - std::vector> remora_runners; + std::vector remora_runners; std::vector modbase_devices; int remora_callers = 1; diff --git a/dorado/nn/Runners.h b/dorado/nn/Runners.h index 9ca6624c..4e1aadc4 100644 --- a/dorado/nn/Runners.h +++ b/dorado/nn/Runners.h @@ -16,7 +16,8 @@ using Runner = std::shared_ptr; namespace modbase { class ModBaseRunner; -} +using Runner = std::unique_ptr; +} // namespace modbase std::pair, size_t> create_basecall_runners( const dorado::basecall::CRFModelConfig& model_config, @@ -28,7 +29,7 @@ std::pair, size_t> create_basecall_runners float memory_fraction, bool guard_gpus); -std::vector> create_modbase_runners( +std::vector create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, diff --git a/dorado/read_pipeline/ModBaseCallerNode.cpp b/dorado/read_pipeline/ModBaseCallerNode.cpp index e5dfd58b..57e0912b 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.cpp +++ b/dorado/read_pipeline/ModBaseCallerNode.cpp @@ -50,11 +50,10 @@ struct ModBaseCallerNode::WorkingRead { num_modbase_chunks_called; // Number of modbase chunks which have been scored }; -ModBaseCallerNode::ModBaseCallerNode( - std::vector> model_runners, - size_t remora_threads, - size_t block_stride, - size_t max_reads) +ModBaseCallerNode::ModBaseCallerNode(std::vector model_runners, + size_t remora_threads, + size_t block_stride, + size_t max_reads) : MessageSink(max_reads), m_runners(std::move(model_runners)), m_num_input_workers(remora_threads), diff --git a/dorado/read_pipeline/ModBaseCallerNode.h b/dorado/read_pipeline/ModBaseCallerNode.h index 968c0523..cca1660b 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.h +++ b/dorado/read_pipeline/ModBaseCallerNode.h @@ -18,13 +18,15 @@ namespace dorado { namespace modbase { class ModBaseRunner; -} +using Runner = std::unique_ptr; +} // namespace modbase + class ModBaseCallerNode : public MessageSink { struct RemoraChunk; struct WorkingRead; public: - ModBaseCallerNode(std::vector> model_runners, + ModBaseCallerNode(std::vector model_runners, size_t remora_threads, size_t block_stride, size_t max_reads); @@ -57,7 +59,7 @@ class ModBaseCallerNode : public MessageSink { // Worker thread, processes chunk results back into the reads void output_worker_thread(); - std::vector> m_runners; + std::vector m_runners; size_t m_num_input_workers = 0; size_t m_block_stride; size_t m_batch_size; diff --git a/dorado/read_pipeline/Pipelines.cpp b/dorado/read_pipeline/Pipelines.cpp index bb19a51f..c20125fe 100644 --- a/dorado/read_pipeline/Pipelines.cpp +++ b/dorado/read_pipeline/Pipelines.cpp @@ -16,18 +16,17 @@ namespace dorado::pipelines { -void create_simplex_pipeline( - PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector>&& modbase_runners, - size_t overlap, - uint32_t mean_qscore_start_pos, - int scaler_node_threads, - bool enable_read_splitter, - int splitter_node_threads, - int modbase_node_threads, - NodeHandle sink_node_handle, - NodeHandle source_node_handle) { +void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, + std::vector&& runners, + std::vector&& modbase_runners, + size_t overlap, + uint32_t mean_qscore_start_pos, + int scaler_node_threads, + bool enable_read_splitter, + int splitter_node_threads, + int modbase_node_threads, + NodeHandle sink_node_handle, + NodeHandle source_node_handle) { const auto& model_config = runners.front()->config(); auto model_stride = runners.front()->model_stride(); auto adjusted_overlap = (overlap / model_stride) * model_stride; @@ -106,19 +105,18 @@ void create_simplex_pipeline( } } -void create_stereo_duplex_pipeline( - PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, - std::vector>&& modbase_runners, - size_t overlap, - uint32_t mean_qscore_start_pos, - int scaler_node_threads, - int splitter_node_threads, - int modbase_node_threads, - PairingParameters pairing_parameters, - NodeHandle sink_node_handle, - NodeHandle source_node_handle) { +void create_stereo_duplex_pipeline(PipelineDescriptor& pipeline_desc, + std::vector&& runners, + std::vector&& stereo_runners, + std::vector&& modbase_runners, + size_t overlap, + uint32_t mean_qscore_start_pos, + int scaler_node_threads, + int splitter_node_threads, + int modbase_node_threads, + PairingParameters pairing_parameters, + NodeHandle sink_node_handle, + NodeHandle source_node_handle) { const auto& model_config = runners.front()->config(); const auto& stereo_model_config = stereo_runners.front()->config(); std::string model_name = diff --git a/dorado/read_pipeline/Pipelines.h b/dorado/read_pipeline/Pipelines.h index d68dea4a..df53cd8f 100644 --- a/dorado/read_pipeline/Pipelines.h +++ b/dorado/read_pipeline/Pipelines.h @@ -18,7 +18,8 @@ using Runner = std::shared_ptr; namespace modbase { class ModBaseRunner; -} +using Runner = std::unique_ptr; +} // namespace modbase using PairingParameters = std::variant>; @@ -27,35 +28,33 @@ namespace pipelines { /// Create a simplex basecall pipeline description /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline -void create_simplex_pipeline( - PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector>&& modbase_runners, - size_t overlap, - uint32_t mean_qscore_start_pos, - int scaler_node_threads, - bool enable_read_splitter, - int splitter_node_threads, - int modbase_threads, - NodeHandle sink_node_handle, - NodeHandle source_node_handle); +void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, + std::vector&& runners, + std::vector&& modbase_runners, + size_t overlap, + uint32_t mean_qscore_start_pos, + int scaler_node_threads, + bool enable_read_splitter, + int splitter_node_threads, + int modbase_threads, + NodeHandle sink_node_handle, + NodeHandle source_node_handle); /// Create a duplex basecall pipeline description /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline -void create_stereo_duplex_pipeline( - PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, - std::vector>&& modbase_runners, - size_t overlap, - uint32_t mean_qscore_start_pos, - int scaler_node_threads, - int splitter_node_threads, - int modbase_node_threads, - PairingParameters pairing_parameters, - NodeHandle sink_node_handle, - NodeHandle source_node_handle); +void create_stereo_duplex_pipeline(PipelineDescriptor& pipeline_desc, + std::vector&& runners, + std::vector&& stereo_runners, + std::vector&& modbase_runners, + size_t overlap, + uint32_t mean_qscore_start_pos, + int scaler_node_threads, + int splitter_node_threads, + int modbase_node_threads, + PairingParameters pairing_parameters, + NodeHandle sink_node_handle, + NodeHandle source_node_handle); } // namespace pipelines diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index dadac3e6..560263b5 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -276,7 +276,7 @@ DEFINE_TEST(NodeSmokeTestRead, "ModBaseCallerNode") { dorado::basecall::load_crf_model_config(model_dir.m_path / model_name).stride; // Create runners - std::vector> remora_runners; + std::vector remora_runners; std::vector modbase_devices; int batch_size = default_params.remora_batchsize; if (gpu) { From eb4780f8db8b2c71931c2cd66278f11b427fd2ee Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 15:05:56 +0000 Subject: [PATCH 07/23] Basecall model runner to unique_ptr --- dorado/basecall/ModelRunner.h | 2 +- dorado/cli/basecaller.cpp | 1 - dorado/cli/duplex.cpp | 1 - dorado/nn/Runners.cpp | 13 +++++-------- dorado/nn/Runners.h | 10 +++------- dorado/read_pipeline/BasecallerNode.cpp | 4 +++- dorado/read_pipeline/BasecallerNode.h | 6 +++++- dorado/read_pipeline/Pipelines.h | 2 +- tests/NodeSmokeTest.cpp | 6 +++--- 9 files changed, 21 insertions(+), 24 deletions(-) diff --git a/dorado/basecall/ModelRunner.h b/dorado/basecall/ModelRunner.h index 6b43b0e2..f2e65078 100644 --- a/dorado/basecall/ModelRunner.h +++ b/dorado/basecall/ModelRunner.h @@ -26,7 +26,7 @@ class ModelRunnerBase { virtual stats::NamedStats sample_stats() const = 0; }; -using Runner = std::shared_ptr; +using Runner = std::unique_ptr; template class ModelRunner final : public ModelRunnerBase { diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp index 5a57a01a..8608fa46 100644 --- a/dorado/cli/basecaller.cpp +++ b/dorado/cli/basecaller.cpp @@ -3,7 +3,6 @@ #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" -#include "modbase/ModBaseRunner.h" #include "models/kits.h" #include "models/models.h" #include "nn/Runners.h" diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index 123a7d43..2829813c 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -3,7 +3,6 @@ #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" -#include "modbase/ModBaseRunner.h" #include "models/kits.h" #include "models/metadata.h" #include "models/models.h" diff --git a/dorado/nn/Runners.cpp b/dorado/nn/Runners.cpp index c3e5755c..6764a5a1 100644 --- a/dorado/nn/Runners.cpp +++ b/dorado/nn/Runners.cpp @@ -1,10 +1,7 @@ #include "Runners.h" -#include "basecall/CRFModel.h" -#include "basecall/ModelRunner.h" #include "basecall/crf_utils.h" #include "basecall/decode/CPUDecoder.h" -#include "modbase/ModBaseRunner.h" #if DORADO_GPU_BUILD #ifdef __APPLE__ @@ -56,7 +53,7 @@ std::pair, size_t> create_basecall_runners num_cpu_runners); for (size_t i = 0; i < num_cpu_runners; i++) { - runners.push_back(std::make_shared>( + runners.push_back(std::make_unique>( model_config, device, int(chunk_size), int(batch_size))); } } @@ -65,7 +62,7 @@ std::pair, size_t> create_basecall_runners else if (device == "metal") { auto caller = basecall::create_metal_caller(model_config, chunk_size, batch_size); for (size_t i = 0; i < num_gpu_runners; i++) { - runners.push_back(std::make_shared(caller)); + runners.push_back(std::make_unique(caller)); } if (batch_size == 0) { spdlog::info(" - set batch size to {}", runners.back()->batch_size()); @@ -101,7 +98,7 @@ std::pair, size_t> create_basecall_runners for (size_t j = 0; j < num_devices; j++) { for (size_t i = 0; i < num_gpu_runners; i++) { - runners.push_back(std::make_shared(callers[j])); + runners.push_back(std::make_unique(callers[j])); } if (batch_size == 0) { spdlog::info(" - set batch size for {} to {}", devices[j], @@ -123,7 +120,7 @@ std::pair, size_t> create_basecall_runners auto model_stride = runners.front()->model_stride(); #endif auto adjusted_chunk_size = runners.front()->chunk_size(); - assert(std::all_of(runners.begin(), runners.end(), [&](auto runner) { + assert(std::all_of(runners.begin(), runners.end(), [&](const auto& runner) { return runner->model_stride() == model_stride && runner->chunk_size() == adjusted_chunk_size; })); @@ -134,7 +131,7 @@ std::pair, size_t> create_basecall_runners chunk_size = adjusted_chunk_size; } - return {runners, num_devices}; + return {std::move(runners), num_devices}; } std::vector create_modbase_runners( diff --git a/dorado/nn/Runners.h b/dorado/nn/Runners.h index 4e1aadc4..98b201df 100644 --- a/dorado/nn/Runners.h +++ b/dorado/nn/Runners.h @@ -1,5 +1,8 @@ #pragma once +#include "basecall/ModelRunner.h" +#include "modbase/ModBaseRunner.h" + #include #include #include @@ -10,15 +13,8 @@ namespace dorado { namespace basecall { struct CRFModelConfig; -class ModelRunnerBase; -using Runner = std::shared_ptr; } // namespace basecall -namespace modbase { -class ModBaseRunner; -using Runner = std::unique_ptr; -} // namespace modbase - std::pair, size_t> create_basecall_runners( const dorado::basecall::CRFModelConfig& model_config, const std::string& device, diff --git a/dorado/read_pipeline/BasecallerNode.cpp b/dorado/read_pipeline/BasecallerNode.cpp index 6d8f72fa..5e200890 100644 --- a/dorado/read_pipeline/BasecallerNode.cpp +++ b/dorado/read_pipeline/BasecallerNode.cpp @@ -1,10 +1,12 @@ #include "BasecallerNode.h" +#include "basecall/ModelRunner.h" #include "basecall/decode/CPUDecoder.h" #include "stitch.h" #include "utils/stats.h" #include +#include #include #include @@ -116,7 +118,7 @@ void BasecallerNode::input_worker_thread() { void BasecallerNode::basecall_current_batch(int worker_id) { NVTX3_FUNC_RANGE(); - auto model_runner = m_model_runners[worker_id]; + auto &model_runner = m_model_runners[worker_id]; dorado::stats::Timer timer; auto decode_results = model_runner->call_chunks(int(m_batched_chunks[worker_id].size())); m_call_chunks_ms += timer.GetElapsedMS(); diff --git a/dorado/read_pipeline/BasecallerNode.h b/dorado/read_pipeline/BasecallerNode.h index f99dc5f1..cec0e1c5 100644 --- a/dorado/read_pipeline/BasecallerNode.h +++ b/dorado/read_pipeline/BasecallerNode.h @@ -1,7 +1,6 @@ #pragma once #include "ReadPipeline.h" -#include "basecall/ModelRunner.h" #include "utils/AsyncQueue.h" #include "utils/stats.h" @@ -14,6 +13,11 @@ namespace dorado { +namespace basecall { +class ModelRunnerBase; +using Runner = std::unique_ptr; +} // namespace basecall + class BasecallerNode : public MessageSink { struct BasecallingRead; struct BasecallingChunk; diff --git a/dorado/read_pipeline/Pipelines.h b/dorado/read_pipeline/Pipelines.h index df53cd8f..233edfd8 100644 --- a/dorado/read_pipeline/Pipelines.h +++ b/dorado/read_pipeline/Pipelines.h @@ -13,7 +13,7 @@ namespace dorado { namespace basecall { class ModelRunnerBase; -using Runner = std::shared_ptr; +using Runner = std::unique_ptr; } // namespace basecall namespace modbase { diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 560263b5..1123de65 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -216,7 +216,7 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { auto caller = dorado::basecall::create_metal_caller(model_config, default_params.chunksize, batch_size); for (int i = 0; i < default_params.num_runners; i++) { - runners.push_back(std::make_shared(caller)); + runners.push_back(std::make_unique(caller)); } #else // __APPLE__ auto devices = dorado::utils::parse_cuda_device_string("cuda:all"); @@ -227,7 +227,7 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { auto caller = dorado::basecall::create_cuda_caller( model_config, default_params.chunksize, int(batch_size), device, 1.f, true); for (int i = 0; i < default_params.num_runners; i++) { - runners.push_back(std::make_shared(caller)); + runners.push_back(std::make_unique(caller)); } } #endif // __APPLE__ @@ -239,7 +239,7 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { set_num_reads(5); set_expected_messages(5); batch_size = 8; - runners.push_back(std::make_shared>( + runners.push_back(std::make_unique>( model_config, "cpu", default_params.chunksize, int(batch_size))); } From 17d7597c4e3a35ecb0dd1b480318c49145e9ee10 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 15:13:04 +0000 Subject: [PATCH 08/23] Move generic top level functions to their own folder --- CMakeLists.txt | 8 ++++---- dorado/{read_pipeline => api}/Pipelines.cpp | 12 ++++++------ dorado/{read_pipeline => api}/Pipelines.h | 2 +- dorado/{nn => api}/Runners.cpp | 0 dorado/{nn => api}/Runners.h | 0 dorado/cli/basecaller.cpp | 4 ++-- dorado/cli/duplex.cpp | 4 ++-- 7 files changed, 15 insertions(+), 15 deletions(-) rename dorado/{read_pipeline => api}/Pipelines.cpp (97%) rename dorado/{read_pipeline => api}/Pipelines.h (98%) rename dorado/{nn => api}/Runners.cpp (100%) rename dorado/{nn => api}/Runners.h (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index e33aa3e2..cf1804ec 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -181,10 +181,10 @@ set(LIB_SOURCE_FILES dorado/alignment/Minimap2Index.h dorado/alignment/Minimap2IndexSupportTypes.h dorado/alignment/Minimap2Options.h - dorado/nn/Runners.cpp - dorado/nn/Runners.h - dorado/read_pipeline/Pipelines.cpp - dorado/read_pipeline/Pipelines.h + dorado/api/Runners.cpp + dorado/api/Runners.h + dorado/api/Pipelines.cpp + dorado/api/Pipelines.h dorado/read_pipeline/FakeDataLoader.cpp dorado/read_pipeline/FakeDataLoader.h dorado/read_pipeline/ReadPipeline.cpp diff --git a/dorado/read_pipeline/Pipelines.cpp b/dorado/api/Pipelines.cpp similarity index 97% rename from dorado/read_pipeline/Pipelines.cpp rename to dorado/api/Pipelines.cpp index c20125fe..8c286bb8 100644 --- a/dorado/read_pipeline/Pipelines.cpp +++ b/dorado/api/Pipelines.cpp @@ -1,14 +1,14 @@ #include "Pipelines.h" -#include "BasecallerNode.h" -#include "ModBaseCallerNode.h" -#include "PairingNode.h" -#include "ReadSplitNode.h" -#include "ScalerNode.h" -#include "StereoDuplexEncoderNode.h" #include "basecall/CRFModelConfig.h" #include "basecall/ModelRunner.h" #include "modbase/ModBaseRunner.h" +#include "read_pipeline/BasecallerNode.h" +#include "read_pipeline/ModBaseCallerNode.h" +#include "read_pipeline/PairingNode.h" +#include "read_pipeline/ReadSplitNode.h" +#include "read_pipeline/ScalerNode.h" +#include "read_pipeline/StereoDuplexEncoderNode.h" #include "splitter/DuplexReadSplitter.h" #include "splitter/RNAReadSplitter.h" diff --git a/dorado/read_pipeline/Pipelines.h b/dorado/api/Pipelines.h similarity index 98% rename from dorado/read_pipeline/Pipelines.h rename to dorado/api/Pipelines.h index 233edfd8..d4b42cba 100644 --- a/dorado/read_pipeline/Pipelines.h +++ b/dorado/api/Pipelines.h @@ -1,6 +1,6 @@ #pragma once -#include "ReadPipeline.h" +#include "read_pipeline/ReadPipeline.h" #include #include diff --git a/dorado/nn/Runners.cpp b/dorado/api/Runners.cpp similarity index 100% rename from dorado/nn/Runners.cpp rename to dorado/api/Runners.cpp diff --git a/dorado/nn/Runners.h b/dorado/api/Runners.h similarity index 100% rename from dorado/nn/Runners.h rename to dorado/api/Runners.h diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp index 8608fa46..ee1e88ae 100644 --- a/dorado/cli/basecaller.cpp +++ b/dorado/cli/basecaller.cpp @@ -1,17 +1,17 @@ #include "Version.h" +#include "api/Pipelines.h" +#include "api/Runners.h" #include "basecall/CRFModelConfig.h" #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" #include "data_loader/ModelFinder.h" #include "models/kits.h" #include "models/models.h" -#include "nn/Runners.h" #include "read_pipeline/AdapterDetectorNode.h" #include "read_pipeline/AlignerNode.h" #include "read_pipeline/BarcodeClassifierNode.h" #include "read_pipeline/HtsReader.h" #include "read_pipeline/HtsWriter.h" -#include "read_pipeline/Pipelines.h" #include "read_pipeline/PolyACalculator.h" #include "read_pipeline/ProgressTracker.h" #include "read_pipeline/ReadFilterNode.h" diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index 2829813c..a1ad6c1e 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -1,4 +1,6 @@ #include "Version.h" +#include "api/Pipelines.h" +#include "api/Runners.h" #include "basecall/CRFModelConfig.h" #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" @@ -6,12 +8,10 @@ #include "models/kits.h" #include "models/metadata.h" #include "models/models.h" -#include "nn/Runners.h" #include "read_pipeline/AlignerNode.h" #include "read_pipeline/BaseSpaceDuplexCallerNode.h" #include "read_pipeline/DuplexReadTaggingNode.h" #include "read_pipeline/HtsWriter.h" -#include "read_pipeline/Pipelines.h" #include "read_pipeline/ProgressTracker.h" #include "read_pipeline/ReadFilterNode.h" #include "read_pipeline/ReadToBamTypeNode.h" From 80f32903246833a0642bb17127484ecd29e66590 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 15:25:33 +0000 Subject: [PATCH 09/23] Add explicit virtual destructor to ModelRunnerBase --- dorado/basecall/ModelRunner.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dorado/basecall/ModelRunner.h b/dorado/basecall/ModelRunner.h index f2e65078..4157da81 100644 --- a/dorado/basecall/ModelRunner.h +++ b/dorado/basecall/ModelRunner.h @@ -14,6 +14,7 @@ namespace dorado::basecall { class ModelRunnerBase { public: + virtual ~ModelRunnerBase() = 0; virtual void accept_chunk(int chunk_idx, const at::Tensor &chunk) = 0; virtual std::vector call_chunks(int num_chunks) = 0; virtual const CRFModelConfig &config() const = 0; @@ -26,6 +27,8 @@ class ModelRunnerBase { virtual stats::NamedStats sample_stats() const = 0; }; +inline ModelRunnerBase::~ModelRunnerBase() = default; + using Runner = std::unique_ptr; template From e98aa13046b9dcc56c10cea9bc464b1ffa067cfa Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 15:43:44 +0000 Subject: [PATCH 10/23] Reuse precompiled headers to avoid windows error --- dorado/basecall/CMakeLists.txt | 4 ++++ dorado/modbase/CMakeLists.txt | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/dorado/basecall/CMakeLists.txt b/dorado/basecall/CMakeLists.txt index 300fe402..bc14010d 100644 --- a/dorado/basecall/CMakeLists.txt +++ b/dorado/basecall/CMakeLists.txt @@ -60,3 +60,7 @@ target_include_directories(dorado_basecall ) enable_warnings_as_errors(dorado_basecall) + +if (DORADO_ENABLE_PCH) + target_precompile_headers(dorado_basecall REUSE_FROM dorado_utils) +endif() diff --git a/dorado/modbase/CMakeLists.txt b/dorado/modbase/CMakeLists.txt index 3a5d0695..ef4e04e6 100644 --- a/dorado/modbase/CMakeLists.txt +++ b/dorado/modbase/CMakeLists.txt @@ -40,3 +40,7 @@ target_link_libraries(dorado_modbase ) enable_warnings_as_errors(dorado_modbase) + +if (DORADO_ENABLE_PCH) + target_precompile_headers(dorado_modbase REUSE_FROM dorado_utils) +endif() From e0a0ef0aeb0913b92812eb0aaf15c434dba4ec92 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Tue, 12 Dec 2023 16:19:05 +0000 Subject: [PATCH 11/23] Fix windows PCH usage --- dorado/basecall/CMakeLists.txt | 4 ++++ dorado/modbase/CMakeLists.txt | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/dorado/basecall/CMakeLists.txt b/dorado/basecall/CMakeLists.txt index bc14010d..c70f8725 100644 --- a/dorado/basecall/CMakeLists.txt +++ b/dorado/basecall/CMakeLists.txt @@ -63,4 +63,8 @@ enable_warnings_as_errors(dorado_basecall) if (DORADO_ENABLE_PCH) target_precompile_headers(dorado_basecall REUSE_FROM dorado_utils) + # these are defined publicly by minimap2, which we're not using. + # we need to define them here so that the environment matches the one + # used when building the PCH + target_compile_definitions(dorado_basecall PRIVATE PTW32_CLEANUP_C PTW32_STATIC_LIB) endif() diff --git a/dorado/modbase/CMakeLists.txt b/dorado/modbase/CMakeLists.txt index ef4e04e6..94204dff 100644 --- a/dorado/modbase/CMakeLists.txt +++ b/dorado/modbase/CMakeLists.txt @@ -43,4 +43,8 @@ enable_warnings_as_errors(dorado_modbase) if (DORADO_ENABLE_PCH) target_precompile_headers(dorado_modbase REUSE_FROM dorado_utils) + # these are defined publicly by minimap2, which we're not using. + # we need to define them here so that the environment matches the one + # used when building the PCH + target_compile_definitions(dorado_modbase PRIVATE PTW32_CLEANUP_C PTW32_STATIC_LIB) endif() From 383ac994f384744e62595d4acd52469e210ae7dc Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 13:37:33 +0000 Subject: [PATCH 12/23] CR: Inherit macro, target_sources --- dorado/basecall/CMakeLists.txt | 17 ++++++----------- dorado/basecall/CRFModelConfig.h | 2 -- dorado/modbase/CMakeLists.txt | 10 +--------- 3 files changed, 7 insertions(+), 22 deletions(-) diff --git a/dorado/basecall/CMakeLists.txt b/dorado/basecall/CMakeLists.txt index c70f8725..88bf5df9 100644 --- a/dorado/basecall/CMakeLists.txt +++ b/dorado/basecall/CMakeLists.txt @@ -1,4 +1,4 @@ -set(LIB_SOURCE_FILES +add_library(dorado_basecall STATIC crf_utils.h crf_utils.cpp CRFModel.h @@ -14,12 +14,14 @@ set(LIB_SOURCE_FILES if (DORADO_GPU_BUILD) if(APPLE) - list(APPEND LIB_SOURCE_FILES + target_sources(dorado_basecall + PRIVATE MetalCRFModel.h MetalCRFModel.cpp ) else() - list(APPEND LIB_SOURCE_FILES + target_sources(dorado_basecall + PRIVATE CudaCRFModel.h CudaCRFModel.cpp decode/GPUDecoder.cpp @@ -28,13 +30,6 @@ if (DORADO_GPU_BUILD) endif() endif() -add_library(dorado_basecall STATIC ${LIB_SOURCE_FILES}) - -target_compile_definitions(dorado_basecall - PUBLIC - DORADO_GPU_BUILD=$ -) - target_include_directories(dorado_basecall SYSTEM PUBLIC @@ -48,8 +43,8 @@ target_link_libraries(dorado_basecall PUBLIC ${TORCH_LIBRARIES} dorado_utils - dorado_models_lib PRIVATE + dorado_models_lib ${KOI_LIBRARIES} spdlog::spdlog ) diff --git a/dorado/basecall/CRFModelConfig.h b/dorado/basecall/CRFModelConfig.h index 18f92c18..d6165c3f 100644 --- a/dorado/basecall/CRFModelConfig.h +++ b/dorado/basecall/CRFModelConfig.h @@ -1,7 +1,5 @@ #pragma once -#include "utils/math_utils.h" - #include #include #include diff --git a/dorado/modbase/CMakeLists.txt b/dorado/modbase/CMakeLists.txt index 94204dff..27a4e9a9 100644 --- a/dorado/modbase/CMakeLists.txt +++ b/dorado/modbase/CMakeLists.txt @@ -1,4 +1,4 @@ -set(LIB_SOURCE_FILES +add_library(dorado_modbase STATIC ModBaseContext.cpp ModBaseContext.h ModbaseEncoder.cpp @@ -15,13 +15,6 @@ set(LIB_SOURCE_FILES MotifMatcher.h ) -add_library(dorado_modbase STATIC ${LIB_SOURCE_FILES}) - -target_compile_definitions(dorado_modbase - PUBLIC - DORADO_GPU_BUILD=$ -) - target_include_directories(dorado_modbase SYSTEM PUBLIC @@ -30,7 +23,6 @@ target_include_directories(dorado_modbase ${DORADO_3RD_PARTY_SOURCE}/toml11 ) - target_link_libraries(dorado_modbase PUBLIC ${TORCH_LIBRARIES} From 63cd9a5cb073423348baf286b192d154515daefb Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 13:37:33 +0000 Subject: [PATCH 13/23] Internal anon namespace --- dorado/basecall/CRFModel.cpp | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/dorado/basecall/CRFModel.cpp b/dorado/basecall/CRFModel.cpp index c0279e84..0d9e3a7a 100644 --- a/dorado/basecall/CRFModel.cpp +++ b/dorado/basecall/CRFModel.cpp @@ -42,14 +42,17 @@ using namespace torch::nn; namespace F = torch::nn::functional; using Slice = torch::indexing::Slice; -#if USE_KOI +namespace dorado::basecall { -static KoiActivation get_koi_activation(dorado::basecall::Activation act) { - if (act == dorado::basecall::Activation::SWISH) { +namespace { + +#if USE_KOI +KoiActivation get_koi_activation(Activation act) { + if (act == Activation::SWISH) { return KOI_SWISH; - } else if (act == dorado::basecall::Activation::SWISH_CLAMP) { + } else if (act == Activation::SWISH_CLAMP) { return KOI_SWISH_CLAMP; - } else if (act == dorado::basecall::Activation::TANH) { + } else if (act == Activation::TANH) { return KOI_TANH; } else { throw std::logic_error("Unrecognised activation function id."); @@ -91,25 +94,24 @@ static KoiActivation get_koi_activation(dorado::basecall::Activation act) { enum class TensorLayout { NTC, TNC, CUTLASS_TNC_F16, CUTLASS_TNC_I8, CUBLAS_TN2C }; // TODO: These should really be part of Koi -static bool koi_can_use_cutlass() { +bool koi_can_use_cutlass() { cudaDeviceProp *prop = at::cuda::getCurrentDeviceProperties(); return ((prop->major == 8 || prop->major == 9) && prop->minor == 0); } -static bool koi_can_use_quantised_lstm() { +bool koi_can_use_quantised_lstm() { cudaDeviceProp *prop = at::cuda::getCurrentDeviceProperties(); // DP4A is supported on Pascal and later, except for TX2 (sm_62). return (prop->major > 6) || (prop->major == 6 && prop->minor != 2); } -static TensorLayout get_koi_lstm_input_layout(int layer_size, - dorado::basecall::Activation activation) { +TensorLayout get_koi_lstm_input_layout(int layer_size, Activation activation) { TensorLayout layout = TensorLayout::CUBLAS_TN2C; if (koi_can_use_quantised_lstm() && (layer_size == 96 || layer_size == 128)) { layout = TensorLayout::NTC; } else if (koi_can_use_cutlass() && layer_size <= 1024 && layer_size > 128 && (layer_size % 128) == 0) { - layout = (activation == dorado::basecall::Activation::TANH) ? TensorLayout::CUTLASS_TNC_I8 - : TensorLayout::CUTLASS_TNC_F16; + layout = (activation == Activation::TANH) ? TensorLayout::CUTLASS_TNC_I8 + : TensorLayout::CUTLASS_TNC_F16; } // Apply override (Cutlass override can only be applied if conditions are met) @@ -274,15 +276,13 @@ class WorkingMemory { #endif // if USE_KOI -namespace { template ModuleHolder populate_model(Model &&model, const std::filesystem::path &path, const at::TensorOptions &options, bool decomposition, bool linear_layer_bias) { - auto state_dict = - dorado::basecall::load_crf_model_weights(path, decomposition, linear_layer_bias); + auto state_dict = load_crf_model_weights(path, decomposition, linear_layer_bias); model->load_state_dict(state_dict); model->to(options.dtype_opt().value().toScalarType()); model->to(options.device_opt().value()); @@ -292,9 +292,8 @@ ModuleHolder populate_model(Model &&model, auto holder = ModuleHolder(module); return holder; } -} // namespace -namespace dorado::basecall { +} // namespace namespace nn { From e736c2f0a122f2409234ff87642f032004247d94 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 13:37:33 +0000 Subject: [PATCH 14/23] Rename non-class-type files --- CMakeLists.txt | 8 ++++---- dorado/api/{Pipelines.cpp => pipeline_creation.cpp} | 2 +- dorado/api/{Pipelines.h => pipeline_creation.h} | 0 dorado/api/{Runners.cpp => runner_creation.cpp} | 2 +- dorado/api/{Runners.h => runner_creation.h} | 0 dorado/cli/basecaller.cpp | 4 ++-- dorado/cli/duplex.cpp | 4 ++-- 7 files changed, 10 insertions(+), 10 deletions(-) rename dorado/api/{Pipelines.cpp => pipeline_creation.cpp} (99%) rename dorado/api/{Pipelines.h => pipeline_creation.h} (100%) rename dorado/api/{Runners.cpp => runner_creation.cpp} (99%) rename dorado/api/{Runners.h => runner_creation.h} (100%) diff --git a/CMakeLists.txt b/CMakeLists.txt index cf1804ec..262a3488 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -181,10 +181,10 @@ set(LIB_SOURCE_FILES dorado/alignment/Minimap2Index.h dorado/alignment/Minimap2IndexSupportTypes.h dorado/alignment/Minimap2Options.h - dorado/api/Runners.cpp - dorado/api/Runners.h - dorado/api/Pipelines.cpp - dorado/api/Pipelines.h + dorado/api/runner_creation.cpp + dorado/api/runner_creation.h + dorado/api/pipeline_creation.cpp + dorado/api/pipeline_creation.h dorado/read_pipeline/FakeDataLoader.cpp dorado/read_pipeline/FakeDataLoader.h dorado/read_pipeline/ReadPipeline.cpp diff --git a/dorado/api/Pipelines.cpp b/dorado/api/pipeline_creation.cpp similarity index 99% rename from dorado/api/Pipelines.cpp rename to dorado/api/pipeline_creation.cpp index 8c286bb8..53e21fdd 100644 --- a/dorado/api/Pipelines.cpp +++ b/dorado/api/pipeline_creation.cpp @@ -1,4 +1,4 @@ -#include "Pipelines.h" +#include "pipeline_creation.h" #include "basecall/CRFModelConfig.h" #include "basecall/ModelRunner.h" diff --git a/dorado/api/Pipelines.h b/dorado/api/pipeline_creation.h similarity index 100% rename from dorado/api/Pipelines.h rename to dorado/api/pipeline_creation.h diff --git a/dorado/api/Runners.cpp b/dorado/api/runner_creation.cpp similarity index 99% rename from dorado/api/Runners.cpp rename to dorado/api/runner_creation.cpp index 6764a5a1..45f1a120 100644 --- a/dorado/api/Runners.cpp +++ b/dorado/api/runner_creation.cpp @@ -1,4 +1,4 @@ -#include "Runners.h" +#include "runner_creation.h" #include "basecall/crf_utils.h" #include "basecall/decode/CPUDecoder.h" diff --git a/dorado/api/Runners.h b/dorado/api/runner_creation.h similarity index 100% rename from dorado/api/Runners.h rename to dorado/api/runner_creation.h diff --git a/dorado/cli/basecaller.cpp b/dorado/cli/basecaller.cpp index ee1e88ae..7c1bc349 100644 --- a/dorado/cli/basecaller.cpp +++ b/dorado/cli/basecaller.cpp @@ -1,6 +1,6 @@ #include "Version.h" -#include "api/Pipelines.h" -#include "api/Runners.h" +#include "api/pipeline_creation.h" +#include "api/runner_creation.h" #include "basecall/CRFModelConfig.h" #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index a1ad6c1e..41fae7db 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -1,6 +1,6 @@ #include "Version.h" -#include "api/Pipelines.h" -#include "api/Runners.h" +#include "api/pipeline_creation.h" +#include "api/runner_creation.h" #include "basecall/CRFModelConfig.h" #include "cli/cli_utils.h" #include "data_loader/DataLoader.h" From a6a1840f3b3e8b2def556a32d5135f7d4ebc1a59 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 13:37:33 +0000 Subject: [PATCH 15/23] Remove unnecessary namespace usage --- dorado/api/pipeline_creation.cpp | 10 +++++----- dorado/api/pipeline_creation.h | 10 +++++----- dorado/api/runner_creation.cpp | 28 ++++++++++++++-------------- dorado/api/runner_creation.h | 6 +++--- dorado/read_pipeline/ScalerNode.cpp | 6 +++--- 5 files changed, 30 insertions(+), 30 deletions(-) diff --git a/dorado/api/pipeline_creation.cpp b/dorado/api/pipeline_creation.cpp index 53e21fdd..9d77dbe7 100644 --- a/dorado/api/pipeline_creation.cpp +++ b/dorado/api/pipeline_creation.cpp @@ -17,8 +17,8 @@ namespace dorado::pipelines { void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, @@ -106,9 +106,9 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, } void create_stereo_duplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& stereo_runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, diff --git a/dorado/api/pipeline_creation.h b/dorado/api/pipeline_creation.h index d4b42cba..e9514bab 100644 --- a/dorado/api/pipeline_creation.h +++ b/dorado/api/pipeline_creation.h @@ -29,8 +29,8 @@ namespace pipelines { /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, @@ -44,9 +44,9 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline void create_stereo_duplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& stereo_runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, diff --git a/dorado/api/runner_creation.cpp b/dorado/api/runner_creation.cpp index 45f1a120..17801075 100644 --- a/dorado/api/runner_creation.cpp +++ b/dorado/api/runner_creation.cpp @@ -19,8 +19,8 @@ namespace dorado { -std::pair, size_t> create_basecall_runners( - const dorado::basecall::CRFModelConfig& model_config, +std::pair, size_t> create_basecall_runners( + const basecall::CRFModelConfig& model_config, const std::string& device, size_t num_gpu_runners, size_t num_cpu_runners, @@ -32,7 +32,7 @@ std::pair, size_t> create_basecall_runners (void)guard_gpus; #endif - std::vector runners; + std::vector runners; // Default is 1 device. CUDA path may alter this. size_t num_devices = 1; @@ -53,7 +53,7 @@ std::pair, size_t> create_basecall_runners num_cpu_runners); for (size_t i = 0; i < num_cpu_runners; i++) { - runners.push_back(std::make_unique>( + runners.push_back(std::make_unique>( model_config, device, int(chunk_size), int(batch_size))); } } @@ -84,12 +84,12 @@ std::pair, size_t> create_basecall_runners cxxpool::thread_pool pool{num_devices}; std::vector> callers; - std::vector>> futures; + std::vector>> futures; for (auto device_string : devices) { - futures.push_back(pool.push(dorado::basecall::create_cuda_caller, model_config, - int(chunk_size), int(batch_size), device_string, - memory_fraction, guard_gpus)); + futures.push_back(pool.push(basecall::create_cuda_caller, model_config, int(chunk_size), + int(batch_size), device_string, memory_fraction, + guard_gpus)); } for (auto& caller : futures) { @@ -98,7 +98,7 @@ std::pair, size_t> create_basecall_runners for (size_t j = 0; j < num_devices; j++) { for (size_t i = 0; i < num_gpu_runners; i++) { - runners.push_back(std::make_unique(callers[j])); + runners.push_back(std::make_unique(callers[j])); } if (batch_size == 0) { spdlog::info(" - set batch size for {} to {}", devices[j], @@ -134,7 +134,7 @@ std::pair, size_t> create_basecall_runners return {std::move(runners), num_devices}; } -std::vector create_modbase_runners( +std::vector create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, @@ -144,7 +144,7 @@ std::vector create_modbase_runners( } // generate model callers before nodes or it affects the speed calculations - std::vector remora_runners; + std::vector remora_runners; std::vector modbase_devices; int remora_callers = 1; @@ -167,10 +167,10 @@ std::vector create_modbase_runners( #endif // DORADO_GPU_BUILD for (const auto& device_string : modbase_devices) { for (int i = 0; i < remora_callers; ++i) { - auto caller = dorado::modbase::create_modbase_caller( - remora_models, int(remora_batch_size), device_string); + auto caller = modbase::create_modbase_caller(remora_models, int(remora_batch_size), + device_string); for (size_t j = 0; j < remora_runners_per_caller; j++) { - remora_runners.push_back(std::make_unique(caller)); + remora_runners.push_back(std::make_unique(caller)); } } }; diff --git a/dorado/api/runner_creation.h b/dorado/api/runner_creation.h index 98b201df..4bb2c2c9 100644 --- a/dorado/api/runner_creation.h +++ b/dorado/api/runner_creation.h @@ -15,8 +15,8 @@ namespace basecall { struct CRFModelConfig; } // namespace basecall -std::pair, size_t> create_basecall_runners( - const dorado::basecall::CRFModelConfig& model_config, +std::pair, size_t> create_basecall_runners( + const basecall::CRFModelConfig& model_config, const std::string& device, size_t num_gpu_runners, size_t num_cpu_runners, @@ -25,7 +25,7 @@ std::pair, size_t> create_basecall_runners float memory_fraction, bool guard_gpus); -std::vector create_modbase_runners( +std::vector create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, diff --git a/dorado/read_pipeline/ScalerNode.cpp b/dorado/read_pipeline/ScalerNode.cpp index 9e1a706c..e42eed82 100644 --- a/dorado/read_pipeline/ScalerNode.cpp +++ b/dorado/read_pipeline/ScalerNode.cpp @@ -20,9 +20,9 @@ using namespace std::chrono_literals; using Slice = at::indexing::Slice; namespace dorado { -using SampleType = dorado::basecall::SampleType; -using ScalingStrategy = dorado::basecall::ScalingStrategy; -using SignalNormalisationParams = dorado::basecall::SignalNormalisationParams; +using SampleType = basecall::SampleType; +using ScalingStrategy = basecall::ScalingStrategy; +using SignalNormalisationParams = basecall::SignalNormalisationParams; std::pair ScalerNode::normalisation(const at::Tensor& x) { // Calculate shift and scale factors for normalisation. From 4d779b2ce46a7532cb179db1f9ac75fa02caeb01 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 13:37:33 +0000 Subject: [PATCH 16/23] Add decode namespace --- dorado/api/runner_creation.cpp | 2 +- dorado/basecall/CudaCRFModel.cpp | 24 ++++++++++++------------ dorado/basecall/CudaCRFModel.h | 2 +- dorado/basecall/MetalCRFModel.cpp | 20 +++++++++++--------- dorado/basecall/MetalCRFModel.h | 2 +- dorado/basecall/ModelRunner.h | 10 +++++----- dorado/basecall/decode/CPUDecoder.cpp | 4 ++-- dorado/basecall/decode/CPUDecoder.h | 4 ++-- dorado/basecall/decode/Decoder.h | 4 ++-- dorado/basecall/decode/GPUDecoder.cpp | 4 ++-- dorado/basecall/decode/GPUDecoder.h | 4 ++-- dorado/basecall/decode/beam_search.cpp | 4 ++++ dorado/basecall/decode/beam_search.h | 2 ++ tests/NodeSmokeTest.cpp | 5 +++-- 14 files changed, 50 insertions(+), 41 deletions(-) diff --git a/dorado/api/runner_creation.cpp b/dorado/api/runner_creation.cpp index 17801075..9e5b91c8 100644 --- a/dorado/api/runner_creation.cpp +++ b/dorado/api/runner_creation.cpp @@ -53,7 +53,7 @@ std::pair, size_t> create_basecall_runners( num_cpu_runners); for (size_t i = 0; i < num_cpu_runners; i++) { - runners.push_back(std::make_unique>( + runners.push_back(std::make_unique>( model_config, device, int(chunk_size), int(batch_size))); } } diff --git a/dorado/basecall/CudaCRFModel.cpp b/dorado/basecall/CudaCRFModel.cpp index 5881d30a..52615639 100644 --- a/dorado/basecall/CudaCRFModel.cpp +++ b/dorado/basecall/CudaCRFModel.cpp @@ -32,16 +32,16 @@ class CudaCaller { : m_config(model_config), m_device(device), m_exclusive_gpu_access(exclusive_gpu_access) { - m_decoder_options = DecoderOptions(); + m_decoder_options = decode::DecoderOptions(); m_decoder_options.q_shift = model_config.qbias; m_decoder_options.q_scale = model_config.qscale; - m_decoder = std::make_unique(model_config.clamp ? 5.f : 0.f); + m_decoder = std::make_unique(model_config.clamp ? 5.f : 0.f); m_num_input_features = model_config.num_features; // adjust chunk size to be a multiple of the stride m_out_chunk_size = chunk_size / model_config.stride; m_in_chunk_size = m_out_chunk_size * model_config.stride; - m_options = at::TensorOptions().dtype(GPUDecoder::dtype).device(device); + m_options = at::TensorOptions().dtype(decode::GPUDecoder::dtype).device(device); assert(m_options.device().is_cuda()); at::InferenceMode guard; @@ -155,7 +155,7 @@ class CudaCaller { // Determine size of working memory for decoder divided by (batch_size * chunk_size) // Decoder needs roughly (beam_width * 4) + num_states + 10 extra bytes // where num_states = 4^(state_len+1) - // See `dorado::GPUDecoder::gpu_part()`, block beginning with `if (!initialized) {` + // See `dorado::decode::GPUDecoder::gpu_part()`, block beginning with `if (!initialized) {` // for more details. int64_t decode_bytes_per_chunk_timestep = 10 + m_decoder_options.beam_width * 4 + (1ull << (model_config.state_len * 2 + 2)); @@ -230,15 +230,15 @@ class CudaCaller { bool done{false}; }; - std::vector call_chunks(at::Tensor &input, - at::Tensor &output, - int num_chunks, - c10::cuda::CUDAStream stream) { + std::vector call_chunks(at::Tensor &input, + at::Tensor &output, + int num_chunks, + c10::cuda::CUDAStream stream) { NVTX3_FUNC_RANGE(); c10::cuda::CUDAStreamGuard stream_guard(stream); if (num_chunks == 0) { - return std::vector(); + return std::vector(); } auto task = std::make_shared(input.to(m_options.device())); @@ -370,8 +370,8 @@ class CudaCaller { const CRFModelConfig m_config; std::string m_device; at::TensorOptions m_options; - std::unique_ptr m_decoder; - DecoderOptions m_decoder_options; + std::unique_ptr m_decoder; + decode::DecoderOptions m_decoder_options; torch::nn::ModuleHolder m_module{nullptr}; std::atomic m_terminate{false}; std::deque> m_input_queue; @@ -413,7 +413,7 @@ void CudaModelRunner::accept_chunk(int chunk_idx, const at::Tensor &chunk) { m_input.index_put_({chunk_idx, torch::indexing::Ellipsis}, chunk); } -std::vector CudaModelRunner::call_chunks(int num_chunks) { +std::vector CudaModelRunner::call_chunks(int num_chunks) { ++m_num_batches_called; stats::Timer timer; auto decoded_chunks = m_caller->call_chunks(m_input, m_output, num_chunks, m_stream); diff --git a/dorado/basecall/CudaCRFModel.h b/dorado/basecall/CudaCRFModel.h index 902f22bd..ce00ebfc 100644 --- a/dorado/basecall/CudaCRFModel.h +++ b/dorado/basecall/CudaCRFModel.h @@ -26,7 +26,7 @@ class CudaModelRunner final : public ModelRunnerBase { public: explicit CudaModelRunner(std::shared_ptr caller); void accept_chunk(int chunk_idx, const at::Tensor& chunk) final; - std::vector call_chunks(int num_chunks) final; + std::vector call_chunks(int num_chunks) final; const CRFModelConfig& config() const final; size_t model_stride() const final; size_t chunk_size() const final; diff --git a/dorado/basecall/MetalCRFModel.cpp b/dorado/basecall/MetalCRFModel.cpp index 502e14bd..890ef28e 100644 --- a/dorado/basecall/MetalCRFModel.cpp +++ b/dorado/basecall/MetalCRFModel.cpp @@ -611,7 +611,7 @@ class MetalCaller { m_device = get_mtl_device(); - m_decoder_options = DecoderOptions(); + m_decoder_options = decode::DecoderOptions(); m_decoder_options.q_shift = model_config.qbias; m_decoder_options.q_scale = model_config.qscale; @@ -815,14 +815,14 @@ class MetalCaller { } struct NNTask { - NNTask(at::Tensor *input_, int num_chunks_, std::vector *out_chunks_) + NNTask(at::Tensor *input_, int num_chunks_, std::vector *out_chunks_) : input(input_), out_chunks(out_chunks_), num_chunks(num_chunks_) {} at::Tensor *input; std::mutex mut; std::condition_variable cv; bool ready{false}; - std::vector *out_chunks; + std::vector *out_chunks; int num_chunks; int decode_chunks_started{0}; int decode_chunks_finished{0}; @@ -830,7 +830,9 @@ class MetalCaller { uint64_t decode_complete_event_id = static_cast(0); }; - void call_chunks(at::Tensor &input, int num_chunks, std::vector &out_chunks) { + void call_chunks(at::Tensor &input, + int num_chunks, + std::vector &out_chunks) { if (num_chunks == 0) { return; } @@ -971,14 +973,14 @@ class MetalCaller { const int out_buf_idx = chunk_idx / m_out_batch_size; const int buf_chunk_idx = chunk_idx % m_out_batch_size; - auto [sequence, qstring, moves] = beam_search_decode( + auto [sequence, qstring, moves] = decode::beam_search_decode( m_scores_int8.at(out_buf_idx).index({Slice(), buf_chunk_idx}), m_bwd.at(out_buf_idx)[buf_chunk_idx], m_posts_int16.at(out_buf_idx)[buf_chunk_idx], m_decoder_options.beam_width, m_decoder_options.beam_cut, m_decoder_options.blank_score, m_decoder_options.q_shift, m_decoder_options.q_scale, score_scale); - (*task->out_chunks)[chunk_idx] = DecodedChunk{sequence, qstring, moves}; + (*task->out_chunks)[chunk_idx] = decode::DecodedChunk{sequence, qstring, moves}; // Wake the waiting thread which called `call_chunks()` if we're done decoding std::unique_lock task_lock(task->mut); @@ -1030,7 +1032,7 @@ class MetalCaller { std::mutex m_decode_lock; std::condition_variable m_decode_cv; std::vector> m_decode_threads; - DecoderOptions m_decoder_options; + decode::DecoderOptions m_decoder_options; nn::MetalModel m_model{nullptr}; NS::SharedPtr m_device; NS::SharedPtr m_bwd_scan_cps, m_fwd_scan_add_softmax_cps; @@ -1077,9 +1079,9 @@ void MetalModelRunner::accept_chunk(int chunk_idx, const at::Tensor &chunk) { } } -std::vector MetalModelRunner::call_chunks(int num_chunks) { +std::vector MetalModelRunner::call_chunks(int num_chunks) { ++m_num_batches_called; - std::vector out_chunks(num_chunks); + std::vector out_chunks(num_chunks); m_caller->call_chunks(m_input, num_chunks, out_chunks); return out_chunks; } diff --git a/dorado/basecall/MetalCRFModel.h b/dorado/basecall/MetalCRFModel.h index e663d855..927faef5 100644 --- a/dorado/basecall/MetalCRFModel.h +++ b/dorado/basecall/MetalCRFModel.h @@ -23,7 +23,7 @@ class MetalModelRunner final : public ModelRunnerBase { public: explicit MetalModelRunner(std::shared_ptr caller); void accept_chunk(int chunk_idx, const at::Tensor& chunk) final; - std::vector call_chunks(int num_chunks) final; + std::vector call_chunks(int num_chunks) final; const CRFModelConfig& config() const final; size_t model_stride() const final; size_t chunk_size() const final; diff --git a/dorado/basecall/ModelRunner.h b/dorado/basecall/ModelRunner.h index 4157da81..70ffd903 100644 --- a/dorado/basecall/ModelRunner.h +++ b/dorado/basecall/ModelRunner.h @@ -16,7 +16,7 @@ class ModelRunnerBase { public: virtual ~ModelRunnerBase() = 0; virtual void accept_chunk(int chunk_idx, const at::Tensor &chunk) = 0; - virtual std::vector call_chunks(int num_chunks) = 0; + virtual std::vector call_chunks(int num_chunks) = 0; virtual const CRFModelConfig &config() const = 0; virtual size_t model_stride() const = 0; virtual size_t chunk_size() const = 0; @@ -39,7 +39,7 @@ class ModelRunner final : public ModelRunnerBase { int chunk_size, int batch_size); void accept_chunk(int chunk_idx, const at::Tensor &chunk) final; - std::vector call_chunks(int num_chunks) final; + std::vector call_chunks(int num_chunks) final; const CRFModelConfig &config() const final { return m_config; }; size_t model_stride() const final { return m_config.stride; } size_t chunk_size() const final { return m_input.size(2); } @@ -54,7 +54,7 @@ class ModelRunner final : public ModelRunnerBase { at::Tensor m_input; at::TensorOptions m_options; std::unique_ptr m_decoder; - DecoderOptions m_decoder_options; + decode::DecoderOptions m_decoder_options; torch::nn::ModuleHolder m_module{nullptr}; // Performance monitoring stats. @@ -69,7 +69,7 @@ ModelRunner::ModelRunner(const CRFModelConfig &model_config, int chunk_size, int batch_size) : m_config(model_config) { - m_decoder_options = DecoderOptions(); + m_decoder_options = decode::DecoderOptions(); m_decoder_options.q_shift = model_config.qbias; m_decoder_options.q_scale = model_config.qscale; m_decoder = std::make_unique(); @@ -85,7 +85,7 @@ ModelRunner::ModelRunner(const CRFModelConfig &model_config, } template -std::vector ModelRunner::call_chunks(int num_chunks) { +std::vector ModelRunner::call_chunks(int num_chunks) { at::InferenceMode guard; dorado::stats::Timer timer; auto scores = m_module->forward(m_input.to(m_options.device_opt().value())); diff --git a/dorado/basecall/decode/CPUDecoder.cpp b/dorado/basecall/decode/CPUDecoder.cpp index f95d37c6..790038aa 100644 --- a/dorado/basecall/decode/CPUDecoder.cpp +++ b/dorado/basecall/decode/CPUDecoder.cpp @@ -86,7 +86,7 @@ at::Tensor backward_scores(const at::Tensor& scores, const float fixed_stay_scor } // namespace -namespace dorado { +namespace dorado::decode { std::vector CPUDecoder::beam_search(const at::Tensor& scores, const int num_chunks, @@ -144,4 +144,4 @@ std::vector CPUDecoder::beam_search(const at::Tensor& scores, return chunk_results; } -} // namespace dorado +} // namespace dorado::decode diff --git a/dorado/basecall/decode/CPUDecoder.h b/dorado/basecall/decode/CPUDecoder.h index b40ad4f7..85fc0ff4 100644 --- a/dorado/basecall/decode/CPUDecoder.h +++ b/dorado/basecall/decode/CPUDecoder.h @@ -4,7 +4,7 @@ #include -namespace dorado { +namespace dorado::decode { class CPUDecoder final : Decoder { public: @@ -14,4 +14,4 @@ class CPUDecoder final : Decoder { constexpr static at::ScalarType dtype = at::ScalarType::Float; }; -} // namespace dorado +} // namespace dorado::decode diff --git a/dorado/basecall/decode/Decoder.h b/dorado/basecall/decode/Decoder.h index 3c354769..9a1b75b5 100644 --- a/dorado/basecall/decode/Decoder.h +++ b/dorado/basecall/decode/Decoder.h @@ -5,7 +5,7 @@ #include #include -namespace dorado { +namespace dorado::decode { struct DecodedChunk { std::string sequence; @@ -30,4 +30,4 @@ class Decoder { const DecoderOptions& options) = 0; }; -} // namespace dorado +} // namespace dorado::decode diff --git a/dorado/basecall/decode/GPUDecoder.cpp b/dorado/basecall/decode/GPUDecoder.cpp index f3fa7766..ee5cebff 100644 --- a/dorado/basecall/decode/GPUDecoder.cpp +++ b/dorado/basecall/decode/GPUDecoder.cpp @@ -11,7 +11,7 @@ extern "C" { #include "koi.h" } -namespace dorado { +namespace dorado::decode { at::Tensor GPUDecoder::gpu_part(at::Tensor scores, DecoderOptions options) { c10::cuda::CUDAGuard device_guard(scores.device()); @@ -113,4 +113,4 @@ std::vector GPUDecoder::beam_search(const at::Tensor &scores, return cpu_part(gpu_part(scores, options)); } -} // namespace dorado +} // namespace dorado::decode diff --git a/dorado/basecall/decode/GPUDecoder.h b/dorado/basecall/decode/GPUDecoder.h index bf7f90bb..57093a0a 100644 --- a/dorado/basecall/decode/GPUDecoder.h +++ b/dorado/basecall/decode/GPUDecoder.h @@ -4,7 +4,7 @@ #include -namespace dorado { +namespace dorado::decode { class GPUDecoder final : Decoder { public: @@ -24,4 +24,4 @@ class GPUDecoder final : Decoder { float m_score_clamp_val; }; -} // namespace dorado +} // namespace dorado::decode diff --git a/dorado/basecall/decode/beam_search.cpp b/dorado/basecall/decode/beam_search.cpp index 7144fa90..38b5c433 100644 --- a/dorado/basecall/decode/beam_search.cpp +++ b/dorado/basecall/decode/beam_search.cpp @@ -117,6 +117,8 @@ uint32_t crc32c(uint32_t crc, uint32_t new_bits) { } // anonymous namespace +namespace dorado::decode { + template float beam_search(const T* const scores, size_t scores_block_stride, @@ -586,3 +588,5 @@ std::tuple> beam_search_decode( return {std::move(sequence), std::move(qstring), std::move(moves)}; } + +} // namespace dorado::decode diff --git a/dorado/basecall/decode/beam_search.h b/dorado/basecall/decode/beam_search.h index 9b38d8f3..80b1c80a 100644 --- a/dorado/basecall/decode/beam_search.h +++ b/dorado/basecall/decode/beam_search.h @@ -8,6 +8,7 @@ #include #include +namespace dorado::decode { std::tuple> beam_search_decode( const at::Tensor& scores_t, const at::Tensor& back_guides_t, @@ -18,3 +19,4 @@ std::tuple> beam_search_decode( float q_shift, float q_scale, float byte_score_scale); +} // namespace dorado::decode diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 1123de65..55472206 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -239,8 +239,9 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { set_num_reads(5); set_expected_messages(5); batch_size = 8; - runners.push_back(std::make_unique>( - model_config, "cpu", default_params.chunksize, int(batch_size))); + runners.push_back( + std::make_unique>( + model_config, "cpu", default_params.chunksize, int(batch_size))); } run_smoke_test(std::move(runners), From 0462e38120d374295a298bdd2dc364c0513a8663 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 14:08:54 +0000 Subject: [PATCH 17/23] Fix bad merge --- dorado/basecall/MetalCRFModel.cpp | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/dorado/basecall/MetalCRFModel.cpp b/dorado/basecall/MetalCRFModel.cpp index 7358d5dc..6215b467 100644 --- a/dorado/basecall/MetalCRFModel.cpp +++ b/dorado/basecall/MetalCRFModel.cpp @@ -4,7 +4,6 @@ #include "MetalCRFModel.h" #include "CRFModelConfig.h" -#include "crf_utils.h" #include "decode/beam_search.h" #include "utils/math_utils.h" #include "utils/metal_utils.h" @@ -40,7 +39,7 @@ constexpr int kTileSize = 8; constexpr int kSIMDGroupWidth = 32; // Returns true on success. -bool finishCommandBuffer(const char *label, MTL::CommandBuffer *cb, int try_count) { +bool finishCommandBuffer(std::string_view label, MTL::CommandBuffer *cb, int try_count) { cb->commit(); cb->waitUntilCompleted(); @@ -490,9 +489,12 @@ struct MetalBlockImpl : Module { if (!finishCommandBuffer("convolutions", command_buffer, try_count)) { return nullptr; } - command_buffer = m_command_queue->commandBuffer(); + std::string lstm_label = "lstm_rnn0"; for (auto &rnn : {rnn1, rnn2, rnn3, rnn4, rnn5}) { + lstm_label.back()++; + command_buffer = m_command_queue->commandBuffer(); + const int kResBufSize = static_cast(dtype_bytes * kernel_simd_groups * 2 * kTileSize * kTileSize); const int kOutBufSize = @@ -506,10 +508,12 @@ struct MetalBlockImpl : Module { tg_buffer_lens, kernel_thread_groups, kernel_simd_groups * kSIMDGroupWidth); } + + if (!finishCommandBuffer(lstm_label, command_buffer, try_count)) { + return nullptr; + } } - if (!finishCommandBuffer("lstm", command_buffer, try_count)) { - return nullptr; - } + command_buffer = m_command_queue->commandBuffer(); // The output buffers of conv/LSTM layers are not used by the decoding, so @@ -682,7 +686,6 @@ class MetalCaller { // Iterate through batch size candidates to find the most efficient one. int best_batch_size = -1; - long long best_us_per_batch_element = std::numeric_limits::max(); for (int batch_size : test_batch_sizes) { spdlog::debug("Trying batch size {}", batch_size); @@ -825,7 +828,7 @@ class MetalCaller { std::mutex mut; std::condition_variable cv; bool ready{false}; - std::vector *out_chunks; + std::vector *out_chunks; int num_chunks; int decode_chunks_started{0}; int decode_chunks_finished{0}; @@ -983,7 +986,8 @@ class MetalCaller { m_decoder_options.beam_cut, m_decoder_options.blank_score, m_decoder_options.q_shift, m_decoder_options.q_scale, score_scale); - (*task->out_chunks)[chunk_idx] = decode::DecodedChunk{sequence, qstring, moves}; + (*task->out_chunks)[chunk_idx] = + decode::DecodedChunk{std::move(sequence), std::move(qstring), std::move(moves)}; // Wake the waiting thread which called `call_chunks()` if we're done decoding std::unique_lock task_lock(task->mut); From 3eb3d5936c152308afc100be6f7b0dddbee85679 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 15:10:52 +0000 Subject: [PATCH 18/23] Fix apple builds --- dorado/basecall/MetalCRFModel.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dorado/basecall/MetalCRFModel.cpp b/dorado/basecall/MetalCRFModel.cpp index 6215b467..34952b1d 100644 --- a/dorado/basecall/MetalCRFModel.cpp +++ b/dorado/basecall/MetalCRFModel.cpp @@ -4,6 +4,7 @@ #include "MetalCRFModel.h" #include "CRFModelConfig.h" +#include "crf_utils.h" #include "decode/beam_search.h" #include "utils/math_utils.h" #include "utils/metal_utils.h" @@ -828,7 +829,7 @@ class MetalCaller { std::mutex mut; std::condition_variable cv; bool ready{false}; - std::vector *out_chunks; + std::vector *out_chunks; int num_chunks; int decode_chunks_started{0}; int decode_chunks_finished{0}; From 8fd06118cf3f66f6520bf55c9a3e27f30c3586ab Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Wed, 13 Dec 2023 17:12:13 +0000 Subject: [PATCH 19/23] Add missing header --- tests/NodeSmokeTest.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index 55472206..b9f71848 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -35,6 +35,10 @@ #include #include +#ifndef _WIN32 +#include +#endif + namespace fs = std::filesystem; namespace { From c766bcb28455ffa8505c6e78e039e8f4612dc2f4 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Thu, 14 Dec 2023 14:34:09 +0000 Subject: [PATCH 20/23] Move decode namespace into basecall namespace --- dorado/api/runner_creation.cpp | 2 +- dorado/basecall/CudaCRFModel.cpp | 2 +- dorado/basecall/decode/CPUDecoder.cpp | 4 ++-- dorado/basecall/decode/CPUDecoder.h | 4 ++-- dorado/basecall/decode/Decoder.h | 4 ++-- dorado/basecall/decode/GPUDecoder.cpp | 4 ++-- dorado/basecall/decode/GPUDecoder.h | 4 ++-- dorado/basecall/decode/beam_search.cpp | 4 ++-- dorado/basecall/decode/beam_search.h | 4 ++-- tests/NodeSmokeTest.cpp | 6 +++--- 10 files changed, 19 insertions(+), 19 deletions(-) diff --git a/dorado/api/runner_creation.cpp b/dorado/api/runner_creation.cpp index 2313d4f3..bbed5fc9 100644 --- a/dorado/api/runner_creation.cpp +++ b/dorado/api/runner_creation.cpp @@ -53,7 +53,7 @@ std::pair, size_t> create_basecall_runners( num_cpu_runners); for (size_t i = 0; i < num_cpu_runners; i++) { - runners.push_back(std::make_unique>( + runners.push_back(std::make_unique>( model_config, device, int(chunk_size), int(batch_size))); } } diff --git a/dorado/basecall/CudaCRFModel.cpp b/dorado/basecall/CudaCRFModel.cpp index 6883a456..09bbc3a7 100644 --- a/dorado/basecall/CudaCRFModel.cpp +++ b/dorado/basecall/CudaCRFModel.cpp @@ -155,7 +155,7 @@ class CudaCaller { // Determine size of working memory for decoder divided by (batch_size * chunk_size) // Decoder needs roughly (beam_width * 4) + num_states + 10 extra bytes // where num_states = 4^(state_len+1) - // See `dorado::decode::GPUDecoder::gpu_part()`, block beginning with `if (!initialized) {` + // See `dorado::basecall::decode::GPUDecoder::gpu_part()`, block beginning with `if (!initialized) {` // for more details. int64_t decode_bytes_per_chunk_timestep = 10 + m_decoder_options.beam_width * 4 + (1ull << (model_config.state_len * 2 + 2)); diff --git a/dorado/basecall/decode/CPUDecoder.cpp b/dorado/basecall/decode/CPUDecoder.cpp index 790038aa..f405c003 100644 --- a/dorado/basecall/decode/CPUDecoder.cpp +++ b/dorado/basecall/decode/CPUDecoder.cpp @@ -86,7 +86,7 @@ at::Tensor backward_scores(const at::Tensor& scores, const float fixed_stay_scor } // namespace -namespace dorado::decode { +namespace dorado::basecall::decode { std::vector CPUDecoder::beam_search(const at::Tensor& scores, const int num_chunks, @@ -144,4 +144,4 @@ std::vector CPUDecoder::beam_search(const at::Tensor& scores, return chunk_results; } -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/dorado/basecall/decode/CPUDecoder.h b/dorado/basecall/decode/CPUDecoder.h index 85fc0ff4..b2011d5c 100644 --- a/dorado/basecall/decode/CPUDecoder.h +++ b/dorado/basecall/decode/CPUDecoder.h @@ -4,7 +4,7 @@ #include -namespace dorado::decode { +namespace dorado::basecall::decode { class CPUDecoder final : Decoder { public: @@ -14,4 +14,4 @@ class CPUDecoder final : Decoder { constexpr static at::ScalarType dtype = at::ScalarType::Float; }; -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/dorado/basecall/decode/Decoder.h b/dorado/basecall/decode/Decoder.h index 9a1b75b5..09ef76bf 100644 --- a/dorado/basecall/decode/Decoder.h +++ b/dorado/basecall/decode/Decoder.h @@ -5,7 +5,7 @@ #include #include -namespace dorado::decode { +namespace dorado::basecall::decode { struct DecodedChunk { std::string sequence; @@ -30,4 +30,4 @@ class Decoder { const DecoderOptions& options) = 0; }; -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/dorado/basecall/decode/GPUDecoder.cpp b/dorado/basecall/decode/GPUDecoder.cpp index ee5cebff..7561371b 100644 --- a/dorado/basecall/decode/GPUDecoder.cpp +++ b/dorado/basecall/decode/GPUDecoder.cpp @@ -11,7 +11,7 @@ extern "C" { #include "koi.h" } -namespace dorado::decode { +namespace dorado::basecall::decode { at::Tensor GPUDecoder::gpu_part(at::Tensor scores, DecoderOptions options) { c10::cuda::CUDAGuard device_guard(scores.device()); @@ -113,4 +113,4 @@ std::vector GPUDecoder::beam_search(const at::Tensor &scores, return cpu_part(gpu_part(scores, options)); } -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/dorado/basecall/decode/GPUDecoder.h b/dorado/basecall/decode/GPUDecoder.h index 57093a0a..46ff4d28 100644 --- a/dorado/basecall/decode/GPUDecoder.h +++ b/dorado/basecall/decode/GPUDecoder.h @@ -4,7 +4,7 @@ #include -namespace dorado::decode { +namespace dorado::basecall::decode { class GPUDecoder final : Decoder { public: @@ -24,4 +24,4 @@ class GPUDecoder final : Decoder { float m_score_clamp_val; }; -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/dorado/basecall/decode/beam_search.cpp b/dorado/basecall/decode/beam_search.cpp index 0cce733e..9dad798b 100644 --- a/dorado/basecall/decode/beam_search.cpp +++ b/dorado/basecall/decode/beam_search.cpp @@ -117,7 +117,7 @@ uint32_t crc32c(uint32_t crc, uint32_t new_bits) { } // anonymous namespace -namespace dorado::decode { +namespace dorado::basecall::decode { template float beam_search(const T* const scores, @@ -589,4 +589,4 @@ std::tuple> beam_search_decode( return {std::move(sequence), std::move(qstring), std::move(moves)}; } -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/dorado/basecall/decode/beam_search.h b/dorado/basecall/decode/beam_search.h index 80b1c80a..5561a280 100644 --- a/dorado/basecall/decode/beam_search.h +++ b/dorado/basecall/decode/beam_search.h @@ -8,7 +8,7 @@ #include #include -namespace dorado::decode { +namespace dorado::basecall::decode { std::tuple> beam_search_decode( const at::Tensor& scores_t, const at::Tensor& back_guides_t, @@ -19,4 +19,4 @@ std::tuple> beam_search_decode( float q_shift, float q_scale, float byte_score_scale); -} // namespace dorado::decode +} // namespace dorado::basecall::decode diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index b9f71848..bee14ec8 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -243,9 +243,9 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { set_num_reads(5); set_expected_messages(5); batch_size = 8; - runners.push_back( - std::make_unique>( - model_config, "cpu", default_params.chunksize, int(batch_size))); + runners.push_back(std::make_unique< + dorado::basecall::ModelRunner>( + model_config, "cpu", default_params.chunksize, int(batch_size))); } run_smoke_test(std::move(runners), From c0d6e58866b292877606fb24c208673de01af022 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Thu, 14 Dec 2023 14:54:43 +0000 Subject: [PATCH 21/23] Rename aliases from `Runner` to `RunnerPtr` to clarify they're smart pointers --- dorado/api/pipeline_creation.cpp | 10 +++++----- dorado/api/pipeline_creation.h | 14 +++++++------- dorado/api/runner_creation.cpp | 8 ++++---- dorado/api/runner_creation.h | 4 ++-- dorado/basecall/ModelRunner.h | 2 +- dorado/cli/duplex.cpp | 2 +- dorado/modbase/ModBaseRunner.h | 2 +- dorado/read_pipeline/BasecallerNode.cpp | 4 ++-- dorado/read_pipeline/BasecallerNode.h | 6 +++--- dorado/read_pipeline/ModBaseCallerNode.cpp | 2 +- dorado/read_pipeline/ModBaseCallerNode.h | 6 +++--- tests/NodeSmokeTest.cpp | 4 ++-- 12 files changed, 32 insertions(+), 32 deletions(-) diff --git a/dorado/api/pipeline_creation.cpp b/dorado/api/pipeline_creation.cpp index 9d77dbe7..b7a20896 100644 --- a/dorado/api/pipeline_creation.cpp +++ b/dorado/api/pipeline_creation.cpp @@ -17,8 +17,8 @@ namespace dorado::pipelines { void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, @@ -106,9 +106,9 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, } void create_stereo_duplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& stereo_runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, diff --git a/dorado/api/pipeline_creation.h b/dorado/api/pipeline_creation.h index e9514bab..ab6dbe45 100644 --- a/dorado/api/pipeline_creation.h +++ b/dorado/api/pipeline_creation.h @@ -13,12 +13,12 @@ namespace dorado { namespace basecall { class ModelRunnerBase; -using Runner = std::unique_ptr; +using RunnerPtr = std::unique_ptr; } // namespace basecall namespace modbase { class ModBaseRunner; -using Runner = std::unique_ptr; +using RunnerPtr = std::unique_ptr; } // namespace modbase using PairingParameters = std::variant>; @@ -29,8 +29,8 @@ namespace pipelines { /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, @@ -44,9 +44,9 @@ void create_simplex_pipeline(PipelineDescriptor& pipeline_desc, /// If source_node_handle is valid, set this to be the source of the simplex pipeline /// If sink_node_handle is valid, set this to be the sink of the simplex pipeline void create_stereo_duplex_pipeline(PipelineDescriptor& pipeline_desc, - std::vector&& runners, - std::vector&& stereo_runners, - std::vector&& modbase_runners, + std::vector&& runners, + std::vector&& stereo_runners, + std::vector&& modbase_runners, size_t overlap, uint32_t mean_qscore_start_pos, int scaler_node_threads, diff --git a/dorado/api/runner_creation.cpp b/dorado/api/runner_creation.cpp index bbed5fc9..80224986 100644 --- a/dorado/api/runner_creation.cpp +++ b/dorado/api/runner_creation.cpp @@ -19,7 +19,7 @@ namespace dorado { -std::pair, size_t> create_basecall_runners( +std::pair, size_t> create_basecall_runners( const basecall::CRFModelConfig& model_config, const std::string& device, size_t num_gpu_runners, @@ -32,7 +32,7 @@ std::pair, size_t> create_basecall_runners( (void)guard_gpus; #endif - std::vector runners; + std::vector runners; // Default is 1 device. CUDA path may alter this. size_t num_devices = 1; @@ -134,7 +134,7 @@ std::pair, size_t> create_basecall_runners( return {std::move(runners), num_devices}; } -std::vector create_modbase_runners( +std::vector create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, @@ -144,7 +144,7 @@ std::vector create_modbase_runners( } // generate model callers before nodes or it affects the speed calculations - std::vector remora_runners; + std::vector remora_runners; std::vector modbase_devices; int remora_callers = 1; diff --git a/dorado/api/runner_creation.h b/dorado/api/runner_creation.h index 4bb2c2c9..b9d09074 100644 --- a/dorado/api/runner_creation.h +++ b/dorado/api/runner_creation.h @@ -15,7 +15,7 @@ namespace basecall { struct CRFModelConfig; } // namespace basecall -std::pair, size_t> create_basecall_runners( +std::pair, size_t> create_basecall_runners( const basecall::CRFModelConfig& model_config, const std::string& device, size_t num_gpu_runners, @@ -25,7 +25,7 @@ std::pair, size_t> create_basecall_runners( float memory_fraction, bool guard_gpus); -std::vector create_modbase_runners( +std::vector create_modbase_runners( const std::vector& remora_models, const std::string& device, size_t remora_runners_per_caller, diff --git a/dorado/basecall/ModelRunner.h b/dorado/basecall/ModelRunner.h index 70ffd903..6ec05b62 100644 --- a/dorado/basecall/ModelRunner.h +++ b/dorado/basecall/ModelRunner.h @@ -29,7 +29,7 @@ class ModelRunnerBase { inline ModelRunnerBase::~ModelRunnerBase() = default; -using Runner = std::unique_ptr; +using RunnerPtr = std::unique_ptr; template class ModelRunner final : public ModelRunnerBase { diff --git a/dorado/cli/duplex.cpp b/dorado/cli/duplex.cpp index 41fae7db..241bd28a 100644 --- a/dorado/cli/duplex.cpp +++ b/dorado/cli/duplex.cpp @@ -487,7 +487,7 @@ int duplex(int argc, char* argv[]) { create_basecall_runners(models.model_config, device, num_runners, 0, batch_size, chunk_size, 0.9f, true); - std::vector stereo_runners; + std::vector stereo_runners; // The fraction argument for GPU memory allocates the fraction of the // _remaining_ memory to the caller. So, we allocate all of the available // memory after simplex caller has been instantiated to the duplex caller. diff --git a/dorado/modbase/ModBaseRunner.h b/dorado/modbase/ModBaseRunner.h index a1e0fa3d..a0bd4501 100644 --- a/dorado/modbase/ModBaseRunner.h +++ b/dorado/modbase/ModBaseRunner.h @@ -49,6 +49,6 @@ class ModBaseRunner { std::atomic m_num_batches_called = 0; }; -using Runner = std::unique_ptr; +using RunnerPtr = std::unique_ptr; } // namespace dorado::modbase diff --git a/dorado/read_pipeline/BasecallerNode.cpp b/dorado/read_pipeline/BasecallerNode.cpp index 5e200890..008a1fc2 100644 --- a/dorado/read_pipeline/BasecallerNode.cpp +++ b/dorado/read_pipeline/BasecallerNode.cpp @@ -290,7 +290,7 @@ void BasecallerNode::basecall_worker_thread(int worker_id) { namespace { // Calculates the input queue size. -size_t CalcMaxChunksIn(const std::vector &model_runners) { +size_t CalcMaxChunksIn(const std::vector &model_runners) { // Allow 2 batches per model runner on the chunks_in queue size_t max_chunks_in = 0; // Allows optimal batch size to be used for every GPU @@ -302,7 +302,7 @@ size_t CalcMaxChunksIn(const std::vector &model_runners) { } // namespace -BasecallerNode::BasecallerNode(std::vector model_runners, +BasecallerNode::BasecallerNode(std::vector model_runners, size_t overlap, int batch_timeout_ms, std::string model_name, diff --git a/dorado/read_pipeline/BasecallerNode.h b/dorado/read_pipeline/BasecallerNode.h index cec0e1c5..5191c2cb 100644 --- a/dorado/read_pipeline/BasecallerNode.h +++ b/dorado/read_pipeline/BasecallerNode.h @@ -15,7 +15,7 @@ namespace dorado { namespace basecall { class ModelRunnerBase; -using Runner = std::unique_ptr; +using RunnerPtr = std::unique_ptr; } // namespace basecall class BasecallerNode : public MessageSink { @@ -24,7 +24,7 @@ class BasecallerNode : public MessageSink { public: // Chunk size and overlap are in raw samples - BasecallerNode(std::vector model_runners, + BasecallerNode(std::vector model_runners, size_t overlap, int batch_timeout_ms, std::string model_name, @@ -50,7 +50,7 @@ class BasecallerNode : public MessageSink { void working_reads_manager(); // Vector of model runners (each with their own GPU access etc) - std::vector m_model_runners; + std::vector m_model_runners; // Chunk length size_t m_chunk_size; // Minimum overlap between two adjacent chunks in a read. Overlap is used to reduce edge effects and improve accuracy. diff --git a/dorado/read_pipeline/ModBaseCallerNode.cpp b/dorado/read_pipeline/ModBaseCallerNode.cpp index 57e0912b..ae56520f 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.cpp +++ b/dorado/read_pipeline/ModBaseCallerNode.cpp @@ -50,7 +50,7 @@ struct ModBaseCallerNode::WorkingRead { num_modbase_chunks_called; // Number of modbase chunks which have been scored }; -ModBaseCallerNode::ModBaseCallerNode(std::vector model_runners, +ModBaseCallerNode::ModBaseCallerNode(std::vector model_runners, size_t remora_threads, size_t block_stride, size_t max_reads) diff --git a/dorado/read_pipeline/ModBaseCallerNode.h b/dorado/read_pipeline/ModBaseCallerNode.h index cca1660b..b3dfe9e4 100644 --- a/dorado/read_pipeline/ModBaseCallerNode.h +++ b/dorado/read_pipeline/ModBaseCallerNode.h @@ -18,7 +18,7 @@ namespace dorado { namespace modbase { class ModBaseRunner; -using Runner = std::unique_ptr; +using RunnerPtr = std::unique_ptr; } // namespace modbase class ModBaseCallerNode : public MessageSink { @@ -26,7 +26,7 @@ class ModBaseCallerNode : public MessageSink { struct WorkingRead; public: - ModBaseCallerNode(std::vector model_runners, + ModBaseCallerNode(std::vector model_runners, size_t remora_threads, size_t block_stride, size_t max_reads); @@ -59,7 +59,7 @@ class ModBaseCallerNode : public MessageSink { // Worker thread, processes chunk results back into the reads void output_worker_thread(); - std::vector m_runners; + std::vector m_runners; size_t m_num_input_workers = 0; size_t m_block_stride; size_t m_batch_size; diff --git a/tests/NodeSmokeTest.cpp b/tests/NodeSmokeTest.cpp index bee14ec8..ef86c218 100644 --- a/tests/NodeSmokeTest.cpp +++ b/tests/NodeSmokeTest.cpp @@ -213,7 +213,7 @@ DEFINE_TEST(NodeSmokeTestRead, "BasecallerNode") { std::size_t batch_size = 128; // Create runners - std::vector runners; + std::vector runners; if (gpu) { #if DORADO_GPU_BUILD #ifdef __APPLE__ @@ -281,7 +281,7 @@ DEFINE_TEST(NodeSmokeTestRead, "ModBaseCallerNode") { dorado::basecall::load_crf_model_config(model_dir.m_path / model_name).stride; // Create runners - std::vector remora_runners; + std::vector remora_runners; std::vector modbase_devices; int batch_size = default_params.remora_batchsize; if (gpu) { From b7548a3b4f93967132e5df295fa99400dc759dd2 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Thu, 14 Dec 2023 14:57:39 +0000 Subject: [PATCH 22/23] Don't need virtual destructor to be pure virtual when we already have pure virtual functions --- dorado/basecall/ModelRunner.h | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/dorado/basecall/ModelRunner.h b/dorado/basecall/ModelRunner.h index 6ec05b62..a9dee625 100644 --- a/dorado/basecall/ModelRunner.h +++ b/dorado/basecall/ModelRunner.h @@ -14,7 +14,7 @@ namespace dorado::basecall { class ModelRunnerBase { public: - virtual ~ModelRunnerBase() = 0; + virtual ~ModelRunnerBase() = default; virtual void accept_chunk(int chunk_idx, const at::Tensor &chunk) = 0; virtual std::vector call_chunks(int num_chunks) = 0; virtual const CRFModelConfig &config() const = 0; @@ -27,8 +27,6 @@ class ModelRunnerBase { virtual stats::NamedStats sample_stats() const = 0; }; -inline ModelRunnerBase::~ModelRunnerBase() = default; - using RunnerPtr = std::unique_ptr; template From d109ab28a6b02cf2b3ab0ade74da217b1f016b57 Mon Sep 17 00:00:00 2001 From: Steve Malton Date: Thu, 14 Dec 2023 15:10:02 +0000 Subject: [PATCH 23/23] Remove unneeded header --- dorado/read_pipeline/BasecallerNode.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/dorado/read_pipeline/BasecallerNode.cpp b/dorado/read_pipeline/BasecallerNode.cpp index 008a1fc2..c39882f6 100644 --- a/dorado/read_pipeline/BasecallerNode.cpp +++ b/dorado/read_pipeline/BasecallerNode.cpp @@ -6,7 +6,6 @@ #include "utils/stats.h" #include -#include #include #include