pirate-frb 0.0.1__tar.gz

pirate_frb-0.0.1/Makefile

@@ -0,0 +1,288 @@
+# This Makefile will be invoked by the python build system (e.g. via 'pip install'),
+# but you can also build individual targets by invoking 'make' directly.
+
+# Disable built-in rules and variables (must be first).
+MAKEFLAGS += --no-builtin-rules 
+MAKEFLAGS += --no-builtin-variables
+
+# Default target 'all' must be first target in Makefile.
+# The 'bin' target builds a bunch of binaries in bin/...
+# The 'lib' target builds the C++ library lib/libpirate.so, and the python extension pirate_frb/pirate_pybind11...so.
+# The 'build_wheel' and 'build_sdist' targets are invoked by 'pip' (or 'make all').
+all: bin lib build_wheel build_sdist
+
+.PHONY: all bin lib build_wheel build_sdist clean
+
+
+####################################################################################################
+#
+# Variables encoding configuration: PYTHON, NVCC, NVCC_ARCH, NVCC_DEPFLAGS.
+#
+# FIXME some day I'll define a configure-script mechanism for setting these variables.
+# For now, if you want to change the defaults, just edit the Makfile.
+
+PYTHON ?= python3
+NVCC ?= nvcc -std=c++17 -m64 -O3 --compiler-options -Wall,-fPIC,-march=x86-64-v3
+
+# Extra nvcc flags needed to build Makefile dependencies
+#   -MMD create dep file, omitting "system" headers
+#   -MP add phony target for each header in dep file (makes error reporting less confusing)
+# Note: we don't need "-MT $@", since we use in-tree object filenames (x.cu -> x.o).
+# Note: we don't need "-MT $*.d", since we use in-tree depfile names (x.cu -> x.d).
+NVCC_DEPFLAGS ?= -MMD -MP
+
+# NVIDIA archictecture.
+DEFAULT_NVCC_ARCH = -gencode arch=compute_80,code=sm_80
+DEFAULT_NVCC_ARCH += -gencode arch=compute_86,code=sm_86
+DEFAULT_NVCC_ARCH += -gencode arch=compute_89,code=sm_89
+# DEFAULT_ARCH += -gencode arch=compute_90,code=sm_90
+NVCC_ARCH ?= $(DEFAULT_NVCC_ARCH)
+
+
+####################################################################################################
+#
+# "Derived" config variables: PYTHON_INCDIR, NUMPY_INCDIR, PYBIND11_INCDIR, PYEXT_SUFFIX, KSGPU_DIR.
+#
+# These are autogenerated by makefile_helper.py, and cached in makefile_helper.out.
+# PYEXT_SUFFIX is something like .cpython-312-x86_64-linux-gnu.so.
+
+
+ifneq ($(MAKECMDGOALS),clean)
+  include makefile_helper.out
+endif
+
+makefile_helper.out: makefile_helper.py Makefile
+	$(PYTHON) makefile_helper.py
+
+
+####################################################################################################
+
+
+# The main output of the build process is these two libraries.
+# Reminder: PYEXT_SUFFIX is something like .cpython-312-x86_64-linux-gnu.so.
+PIRATE_LIB := lib/libpirate.so
+PIRATE_PYEXT := pirate_frb/pirate_pybind11$(PYEXT_SUFFIX)
+
+# These get compiled into lib/libpirate.so.
+LIB_SRCFILES = \
+  src_lib/cpu_downsample.cu \
+  src_lib/file_utils.cu \
+  src_lib/gpu_downsample.cu \
+  src_lib/gpu_transpose.cu \
+  src_lib/system_utils.cu \
+  src_lib/utils.cu \
+  src_lib/DedispersionConfig.cu \
+  src_lib/DedispersionPlan.cu \
+  src_lib/Directory.cu \
+  src_lib/Epoll.cu \
+  src_lib/FakeCorrelator.cu \
+  src_lib/FakeServer.cu \
+  src_lib/File.cu \
+  src_lib/GpuDedispersionKernel.cu \
+  src_lib/GpuLaggedDownsamplingKernel.cu \
+  src_lib/ReferenceDedisperser.cu \
+  src_lib/ReferenceDedispersionKernel.cu \
+  src_lib/ReferenceLagbuf.cu \
+  src_lib/ReferenceLaggedDownsamplingKernel.cu \
+  src_lib/ReferenceTree.cu \
+  src_lib/Socket.cu \
+  src_lib/UntypedArray.cu \
+  src_lib/YamlFile.cu \
+  src_lib/template_instantiations/dedisp_simple_float16.cu \
+  src_lib/template_instantiations/dedisp_simple_float32.cu \
+  src_lib/template_instantiations/dedisp_simple_nolag_float16.cu \
+  src_lib/template_instantiations/dedisp_simple_nolag_float32.cu \
+  src_lib/template_instantiations/dedisp_stage0_float16.cu \
+  src_lib/template_instantiations/dedisp_stage0_float32.cu \
+  src_lib/template_instantiations/dedisp_stage1_float16.cu \
+  src_lib/template_instantiations/dedisp_stage1_float32.cu
+
+# These get compiled into pirate_frb/pirate_pybind11....so.
+PYEXT_SRCFILES = \
+  src_pybind11/pirate_pybind11.cu
+
+# Must list all python source files here.
+# (Otherwise they won't show up in 'pip install' or pypi.)
+PYFILES = \
+  pirate_frb/__init__.py
+
+# These are in 1-1 corresponding with executables in bin/
+# For example, 'src_bin/fake_correlator.cu' gets compiled to 'bin/fake_correlator'.
+BIN_SRCFILES = \
+  src_bin/fake_correlator.cu \
+  src_bin/fake_server.cu \
+  src_bin/scratch.cu \
+  src_bin/show_dedispersion_plan.cu \
+  src_bin/test-avx256-m64-outbuf.cu \
+  src_bin/test-cpu-downsampler.cu \
+  src_bin/test-gpu-dedispersion-kernels.cu \
+  src_bin/test-gpu-downsample.cu \
+  src_bin/test-gpu-lagged-downsampler.cu \
+  src_bin/test-gpu-reduce2.cu \
+  src_bin/test-gpu-transpose.cu \
+  src_bin/test-reference-dedisperser.cu \
+  src_bin/test-reference-tree.cu \
+  src_bin/time-cpu-downsample.cu \
+  src_bin/time-gpu-dedispersion-kernels.cu \
+  src_bin/time-gpu-downsample.cu \
+  src_bin/time-gpu-lagged-downsampler.cu \
+  src_bin/time-gpu-transpose.cu
+
+# Must list all header files here.
+# (Otherwise they won't show up in 'pip install' or pypi.)
+HFILES = \
+  include/pirate/constants.hpp \
+  include/pirate/DedispersionConfig.hpp \
+  include/pirate/DedispersionPlan.hpp \
+  include/pirate/avx256/downsample.hpp \
+  include/pirate/avx256/m64_outbuf.hpp \
+  include/pirate/avx256/m128_outbuf.hpp \
+  include/pirate/gpu/reduce2.hpp \
+  include/pirate/gpu/DownsampleKernel.hpp \
+  include/pirate/gpu/TransposeKernel.hpp \
+  include/pirate/internals/bitvec.hpp \
+  include/pirate/internals/cpu_downsample.hpp \
+  include/pirate/internals/dedispersion_kernel_implementation.hpp \
+  include/pirate/internals/dedispersion_inbufs.hpp \
+  include/pirate/internals/dedispersion_outbufs.hpp \
+  include/pirate/internals/gpu_downsample.hpp \
+  include/pirate/internals/gpu_transpose.hpp \
+  include/pirate/internals/file_utils.hpp \
+  include/pirate/internals/inlines.hpp \
+  include/pirate/internals/system_utils.hpp \
+  include/pirate/internals/utils.hpp \
+  include/pirate/internals/Directory.hpp \
+  include/pirate/internals/Epoll.hpp \
+  include/pirate/internals/FakeCorrelator.hpp \
+  include/pirate/internals/FakeServer.hpp \
+  include/pirate/internals/File.hpp \
+  include/pirate/internals/GpuDedispersionKernel.hpp \
+  include/pirate/internals/ReferenceDedisperser.hpp \
+  include/pirate/internals/ReferenceDedispersionKernel.hpp \
+  include/pirate/internals/ReferenceLagbuf.hpp \
+  include/pirate/internals/ReferenceLaggedDownsamplingKernel.hpp \
+  include/pirate/internals/ReferenceTree.hpp \
+  include/pirate/internals/Socket.hpp \
+  include/pirate/internals/UntypedArray.hpp \
+  include/pirate/internals/YamlFile.hpp
+
+# 'make clean' deletes {*~, *.o, *.d, *.so, *.pyc} from these dirs.
+CLEAN_DIRS := . lib src_bin src_lib src_lib/template_instantiations pirate_frb/__pycache__ include include/pirate include/pirate/avx256 include/pirate/gpu include/pirate/internals
+
+# Extra files to be deleted by 'make clean'.
+# Note that 'pirate_frb/include' and 'pirate_frb/lib' are symlinks, so we put them in CLEAN_FILES, not CLEAN_RMDIRS
+CLEAN_FILES := sdist_files.txt wheel_files.txt makefile_helper.out pirate_frb/include pirate_frb/lib
+
+# Directories that should be empty at the end of 'make clean', and can be deleted.
+CLEAN_RMDIRS := bin lib pirate_frb/__pycache__
+
+
+####################################################################################################
+
+
+LIB_OFILES := $(LIB_SRCFILES:%.cu=%.o)
+PYEXT_OFILES := $(PYEXT_SRCFILES:%.cu=%.o)
+BIN_XFILES := $(BIN_SRCFILES:src_bin/%.cu=bin/%)
+
+# Must include all .d files, or build will break!
+ALL_SRCFILES := $(LIB_SRCFILES) $(PYEXT_SRCFILES) $(BIN_SRCFILES)
+DEPFILES := $(ALL_SRCFILES:%.cu=%.d)
+
+SDIST_FILES := pyproject.toml Makefile makefile_helper.py
+SDIST_FILES += $(PYFILES) $(ALL_SRCFILES) $(HFILES)
+
+# Some symlinks for the wheel:
+#  - header file include/%.hpp gets symlinked to pirate_frb/include/%.hpp
+#  - library lib/libpirate.so gets symlinked to pirate_frb/lib/libpirate.so
+#  - python extension pirate_frb/pirate_pybind11...so does not need to be symlinked/renamed.
+WHEEL_FILES := $(PYFILES) $(PIRATE_PYEXT) pirate_frb/$(PIRATE_LIB)
+WHEEL_FILES += $(HFILES:%=pirate_frb/%)
+
+# Phony targets. The special targets 'build_wheel' and 'build_sdist' are needed by pip/pipmake.
+lib: $(PIRATE_LIB) $(PIRATE_PYEXT)
+bin: $(BIN_XFILES)
+build_wheel: wheel_files.txt $(PIRATE_LIB) $(PIRATE_PYEXT)
+build_sdist: sdist_files.txt
+
+# Symlink {include,lib} into python directory 'pirate_frb'.
+pirate_frb/include:
+	ln -s ../include $@
+pirate_frb/lib:
+	ln -s ../lib $@
+
+# Build object files in src_lib/, src_bin/, and src_lib/template_instantiations/ with default flags.
+%.o: %.cu %.d
+	$(NVCC) $(NVCC_ARCH) $(NVCC_DEPFLAGS) -I$(KSGPU_DIR)/include -c -o $@ $<
+
+# Build object files in src_pybind11/ with special flags.
+src_pybind11/%.o: src_pybind11/%.cu src_pybind11/%.d
+	$(NVCC) $(NVCC_ARCH) $(NVCC_DEPFLAGS) -I$(KSGPU_DIR)/include -I$(PYTHON_INCDIR) -I$(NUMPY_INCDIR) -I$(PYBIND11_INCDIR) -c -o $@ $<
+
+# Build the C++ library (lib/libpirate.so)
+# We want it to automatically pull in the C++ library $(KSGPU_DIR)/lib/libkspgu.so.
+#
+# The python extension has been built correctly if 'objdump -x' shows the following:
+#   NEEDED   libksgpu.so
+#   RUNPATH  $(KSGPU_DIR)/lib    # where Makefile var $(KSGPU_DIR) is read from makefile_helper.out
+#
+# The quoting can be understood by working backwards as follows:
+#  - g++ command line should look like:   g++ -Wl,-rpath="$(KSGPU_DIR)/lib"
+#  - nvcc command line should look like:  nvcc -Xcompiler '"-Wl,-rpath=$(KSGPU_DIR)/lib"'
+
+$(PIRATE_LIB): $(LIB_OFILES)
+	@mkdir -p lib
+	$(NVCC) $(NVCC_ARCH) -shared -o $@ $^ -lksgpu -lyaml-cpp -L$(KSGPU_DIR)/lib -Xcompiler '"-Wl,-rpath=$(KSGPU_DIR)/lib"'
+
+# Build C++ binaries (bin/*)
+# Link flags are similar to previous rule -- see comments above.
+bin/%: src_bin/%.o $(PIRATE_LIB)
+	@mkdir -p bin/
+	$(NVCC) $(NVCC_ARCH) -o $@ $^ -lksgpu -lyaml-cpp -L$(KSGPU_DIR)/lib -Xcompiler '"-Wl,-rpath=$(KSGPU_DIR)/lib"'
+
+# Build the python extension (pirate_frb/pirate_pybind11...so)
+# We want it to automatically pull in the C++ library pirate_frb/lib/libpirate.so.
+#
+# The python extension has been built correctly if 'objdump -x' shows the following:
+#   NEEDED   libpirate.so
+#   RUNPATH  $ORIGIN/lib
+#
+# The quoting can be understood by working backwards as follows:
+#  - g++ command line should look like:   g++ -Wl,-rpath="\$ORIGIN/lib"
+#  - nvcc command line should look like:  nvcc -Xcompiler '"-Wl,-rpath=\\$ORIGIN/lib"'
+#  - Makefile line should look like:      nvcc -Xcompiler '"-Wl,-rpath=\\$$ORIGIN/lib"'
+#
+# Note that we don't link to libksgpu.so or ksgpu_pybind11...so in this step.
+# These libraries end up getting imported as follows:
+#
+#  1. When 'pirate_frb' is imported, we do 'import ksgpu' (in pirate_frb/__init__.py)
+#     before 'import pirate_pybind11'.
+#
+#  2. When 'ksgpu' is imported, we use the "ctypes trick" (see comment in ksgpu/__init__.py)
+#     to load the libraries libksgpu.so and ksgpu_pybind11...so with globally visible symbols.
+
+$(PIRATE_PYEXT): $(PYEXT_OFILES) $(PIRATE_LIB) pirate_frb/lib
+	$(NVCC) $(NVCC_ARCH) -shared -o $@ $(PYEXT_OFILES) -lpirate -Lpirate_frb/lib -Xcompiler '"-Wl,-rpath=\\$$ORIGIN/lib"'
+
+# Needed by pip/pipmake: list of all files that go into the (non-editable) wheel.
+wheel_files.txt: Makefile pirate_frb/include pirate_frb/lib
+	rm -f $@
+	for f in $(WHEEL_FILES); do echo $$f; done >>$@
+
+# Needed by pip/pipmake: list of all files that go into the sdist.
+sdist_files.txt: Makefile
+	rm -f $@
+	for f in $(SDIST_FILES); do echo $$f; done >>$@
+
+clean:
+	@for f in $(foreach d,$(CLEAN_DIRS),$(wildcard $d/*~ $d/*.o $d/*.d $d/*.so $d/*.pyc)); do echo rm $$f; rm $$f; done
+	@for f in $(wildcard $(CLEAN_FILES) $(BIN_XFILES)); do echo rm $$f; rm $$f; done
+	@for d in $(wildcard $(CLEAN_RMDIRS)); do echo rmdir $$d; rmdir $$d; done
+
+# Specifying .SECONDARY with no prerequisites disables auto-deletion of intermediate files.
+.SECONDARY:
+
+# If a depfile is absent, build can still proceed.
+$(DEPFILES):
+
+# Include any depfiles which are present.
+include $(wildcard $(DEPFILES))

pirate_frb-0.0.1/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.2
+Name: pirate_frb
+Version: 0.0.1
+Requires-Python: >=3.8
+Requires-Dist: pybind11
+Requires-Dist: numpy
+Requires-Dist: ksgpu >= 1.0.3

pirate_frb-0.0.1/include/pirate/DedispersionConfig.hpp

@@ -0,0 +1,86 @@
+#ifndef _PIRATE_DEDISPERSION_CONFIG_HPP
+#define _PIRATE_DEDISPERSION_CONFIG_HPP
+
+#include <vector>
+#include <string>
+#include <iostream>
+
+namespace YAML { class Emitter; }      // #include <yaml-cpp/yaml.h>
+namespace pirate { struct YamlFile; }  // #include <pirate/internals/YamlFile.hpp>
+
+
+namespace pirate {
+#if 0
+}  // editor auto-indent
+#endif
+
+
+struct DedispersionConfig
+{
+    // Core dedispersion parameters.
+    ssize_t tree_rank = -1;
+    ssize_t num_downsampling_levels = -1;
+    ssize_t time_samples_per_chunk = 0;
+
+    // For now, there is only one dtype, which can be either "float32" or "float16".
+    // Later, I might split this into "compute" and "ringbuf" dtypes, and allow compressed
+    // dtypes (e.g. float8, int7).
+    
+    std::string dtype;  // "float32" or "float16"
+    
+    struct EarlyTrigger
+    {
+	ssize_t ds_level = -1;
+	ssize_t tree_rank = 0;
+    };
+
+    // Sorted (by ds_level first, then tree_rank).
+    std::vector<EarlyTrigger> early_triggers;
+
+    // GPU configuration.
+    ssize_t beams_per_gpu = 0;
+    ssize_t beams_per_batch = 0;
+    ssize_t num_active_batches = 0;
+    ssize_t gmem_nbytes_per_gpu = 0;
+    
+    void validate() const;
+
+    // Write in informal text format (e.g. for log files)
+    // FIXME I might phase this out, in favor of yaml everywhere.
+    void print(std::ostream &os = std::cout, int indent=0) const;
+
+    // Write in YAML format.
+    void to_yaml(YAML::Emitter &emitter) const;
+    void to_yaml(const std::string &filename) const;    
+    std::string to_yaml_string() const;
+
+    // Construct from YAML file.
+    // The 'verbosity' argument has the following meaning:
+    //   0 = quiet
+    //   1 = announce default values for all unspecified parameters
+    //   2 = announce all parameters
+
+    static DedispersionConfig from_yaml(const std::string &filename, int verbosity=0);
+    static DedispersionConfig from_yaml(const YamlFile &file);
+    
+    // Helper functions for constructing DedispersionConfig instances.
+    // Add early triggers, while maintaining invariant that 'early_triggers' is sorted.
+    void add_early_trigger(ssize_t ds_level, ssize_t tree_rank);
+    void add_early_triggers(ssize_t ds_level, std::initializer_list<ssize_t> tree_ranks);
+
+    // Note: rather than calling this function directly, you probably want the
+    // DedispersionPlan (not DedispersionConfig) member 'nelts_per_segment'.
+    int get_nelts_per_segment() const;
+
+    // make_random(): used for unit tests.
+    static DedispersionConfig make_random();
+};
+
+extern bool operator==(const DedispersionConfig::EarlyTrigger &x, const DedispersionConfig::EarlyTrigger &y);
+extern bool operator>(const DedispersionConfig::EarlyTrigger &x, const DedispersionConfig::EarlyTrigger &y);
+extern std::ostream &operator<<(std::ostream &os, const DedispersionConfig::EarlyTrigger &et);
+
+
+}  // namespace pirate
+
+#endif // _PIRATE_DEDISPERSION_CONFIG_HPP

pirate_frb-0.0.1/include/pirate/DedispersionPlan.hpp

@@ -0,0 +1,100 @@
+#ifndef _PIRATE_DEDISPERSION_PLAN_HPP
+#define _PIRATE_DEDISPERSION_PLAN_HPP
+
+#include "DedispersionConfig.hpp"
+
+#include <vector>
+#include <memory>  // shared_ptr
+#include <ksgpu/Array.hpp>
+
+
+namespace pirate {
+#if 0
+}  // editor auto-indent
+#endif
+
+
+struct DedispersionPlan
+{
+    DedispersionPlan(const DedispersionConfig &config);
+    
+    void print(std::ostream &os=std::cout, int indent=0) const;
+    
+    // --------------------  Helper classes  --------------------
+    
+    struct Stage0Tree
+    {
+	// Note: total tree rank (rank0 + rank1) is equal to (config.tree_rank - (ds_level ? 1 : 0)).
+	
+	int ds_level = -1;  // downsampling level (downsampling "factor" is 2^level)
+	int rank0 = 0;      // rank of Stage0Tree
+	int rank1 = 0;      // rank of subsequent Stage1Tree (if no early trigger)
+	int nt_ds = 0;      // downsampled time samples per chunk (= config.time_samples_per_chunk / pow2(ds_level))
+	
+	int segments_per_beam = 0;   // equal to pow2(rank0+rank1) * (nt_ds / nelts_per_segment)
+        int base_segment = 0;        // cumulative (over all Stage0Trees) segment count
+    };
+
+    struct Stage1Tree
+    {
+	int ds_level = -1;       // Same as Stage0Tree::ds_level
+	int rank0 = 0;           // Same as Stage0Tree::rank0
+	int rank1_ambient = 0;   // Same as Stage0Tree::rank1
+	int rank1_trigger = 0;   // Can be smaller than rank1_ambient, for early trigger
+	int nt_ds = 0;           // Same as Stage0Tree::nt_ds
+		
+	int segments_per_beam = 0;   // equal to pow2(rank0 + rank1_trigger) * (nt_ds / nelts_per_segment)
+        int base_segment = 0;        // cumulative (over all Stage1Trees) segment count
+    };
+
+    struct Ringbuf
+    {
+	long rb_len = 0;           // number of (time chunk, beam) pairs
+	long nseg_per_beam = 0;
+	long base_segment = 0;
+    };
+    
+    // --------------------  Members  --------------------
+    
+    const DedispersionConfig config;
+
+    int nelts_per_segment = 0;   // currently always constants::bytes_per_gpu_cache_line / (sizeof config dtype)
+    int nbytes_per_segment = 0;  // currently always constants::bytes_per_gpu_cache_line
+    
+    std::vector<Stage0Tree> stage0_trees;
+    std::vector<Stage1Tree> stage1_trees;
+
+    ssize_t stage0_total_segments_per_beam = 0;
+    ssize_t stage1_total_segments_per_beam = 0;
+
+    int max_clag = 0;
+    long gmem_ringbuf_nseg = 0;    // includes gmem + g2h + h2g
+
+    // All vector<Ringbuf> objects have length (max_clag + 1).
+    // T = total beams, A = active beams, B = beams per batch.
+    
+    std::vector<Ringbuf> gmem_ringbufs;    // rb_size = (clag*T + A), on GPU
+    std::vector<Ringbuf> g2h_ringbufs;     // rb_size = min(A+B, T), on GPU
+    std::vector<Ringbuf> h2g_ringbufs;     // rb_size = min(A+B, T), on GPU
+    std::vector<Ringbuf> h2h_ringbufs;     // rb_size = (clag*T + B), on host
+
+    // stage0_output_rb_locs, stage1_input_rb_locs.
+    //
+    // These arrays contain GPU ringbuf locations, represented as 4 uint32s:
+    //  uint rb_offset;  // in segments, not bytes
+    //  uint rb_phase;   // index of (time chunk, beam) pair, relative to current pair
+    //  uint rb_len;     // number of (time chunk, beam) pairs in ringbuf (same as Ringbuf::rb_len)
+    //  uint rb_nseg;    // number of segments per (time chunk, beam) (same as Ringbuf::nseg_per_beam)
+    //
+    // The arrays are indexed by:
+    //  iseg0 -> (time/nelts_per_segment, 2^rank1, 2^rank0)
+    //  iseg1 -> (time/nelts_per_segment, 2^rank0, 2^rank1)   note transpose
+
+    ksgpu::Array<uint> stage0_rb_locs;   // shape (stage0_total_segments_per_beam, 4)
+    ksgpu::Array<uint> stage1_rb_locs;   // shape (stage1_total_segments_per_beam, 4)
+};
+
+
+}  // namespace pirate
+
+#endif // _PIRATE_DEDISPERSION_PLAN_HPP