Merge pull request #830 from anarkiwi/tmpl

vkfft support.
IQTLabs · Aug 19, 2023 · c89a344 · c89a344
2 parents 256f20f + b70687a
commit c89a344
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 30 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -58,6 +58,7 @@ RUN apt-get update && apt-get install --no-install-recommends -y -q \
     libspdlog1 \
     libuhd4.1.0 \
     libunwind8 \
+    libvulkan1 \
     libzmq5 \
     mesa-vulkan-drivers \
     python3 \

diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base
@@ -27,13 +27,22 @@ RUN cmake .. && make -j "$(nproc)" && make install
 WORKDIR /root/libsigmf/build
 RUN cmake -DUSE_SYSTEM_JSON=ON -DUSE_SYSTEM_FLATBUFFERS=ON .. && make -j "$(nproc)" && make install
 
+FROM iqtlabs/gnuradio:3.10.7 as vkfft-builder
+ENV DEBIAN_FRONTEND noninteractive
+RUN apt-get update && apt-get install -y --no-install-recommends libvulkan-dev
+WORKDIR /root
+RUN git clone https://github.com/DTolm/VkFFT -b v1.3.1
+WORKDIR /root/VkFFT/build
+RUN cmake .. && make -j "$(nproc)"
+
 FROM iqtlabs/gnuradio:3.10.7 as gr-iqtlabs-builder
 ENV DEBIAN_FRONTEND noninteractive
-COPY --from=sigmf-builder /usr/local /usr/local
 WORKDIR /root
-RUN git clone https://github.com/iqtlabs/gr-iqtlabs -b 1.0.31
+RUN git clone https://github.com/iqtlabs/gr-iqtlabs -b 1.0.32
 WORKDIR /root/gr-iqtlabs/build
-RUN apt-get update && apt-get install -y --no-install-recommends libopencv-dev && cmake .. && make install
+COPY --from=sigmf-builder /usr/local /usr/local
+COPY --from=vkfft-builder /root/VkFFT /root/gr-iqtlabs/VkFFT
+RUN apt-get update && apt-get install -y --no-install-recommends libopencv-dev libvulkan-dev && cmake .. && make -j "$(nproc)" && make install
 
 FROM iqtlabs/gnuradio:3.10.7 as uhd_sample_recorder-builder
 ENV DEBIAN_FRONTEND noninteractive
@@ -52,7 +61,9 @@ COPY --from=uhd_sample_recorder-builder /usr/local /usr/local
 RUN apt-get update && apt-get install -y --no-install-recommends \
     libopencv-core4.5d \
     libopencv-imgcodecs4.5d \
-    libopencv-imgproc4.5d && apt-get -y -q clean && rm -rf /var/lib/apt/lists/*
+    libopencv-imgproc4.5d \
+    libvulkan1 && \
+    apt-get -y -q clean && rm -rf /var/lib/apt/lists/*
 RUN ldconfig -v
 RUN ln -sf /usr/local/lib/python3/dist-packages/* /usr/local/lib/python3.10/dist-packages
 RUN python3 -c "from gnuradio import soapy, iqtlabs"
diff --git a/gamutrf/grscan.py b/gamutrf/grscan.py
@@ -53,6 +53,8 @@ def __init__(
         db_clamp_floor=-200,
         db_clamp_ceil=50,
         rotate_secs=0,
+        fft_batch_size=256,
+        use_vkfft=False,
         sigmf=True,
         iqtlabs=None,
         wavelearner=None,
@@ -67,6 +69,7 @@ def __init__(
         self.sweep_sec = sweep_sec
         self.fft_size = fft_size
         self.wavelearner = wavelearner
+        self.iqtlabs = iqtlabs
         self.samp_rate = samp_rate
 
         ##################################################
@@ -83,8 +86,8 @@ def __init__(
             sdrargs=sdrargs,
         )
 
-        fft_blocks, fft_roll = self.get_fft_blocks(
-            fft_size, dc_block_len, dc_block_long
+        fft_blocks = self.get_fft_blocks(
+            use_vkfft, fft_batch_size, fft_size, dc_block_len, dc_block_long
         )
         self.fft_blocks = fft_blocks + self.get_db_blocks(fft_size, samp_rate, scaling)
         self.fft_to_inference_block = self.fft_blocks[-1]
@@ -96,7 +99,7 @@ def __init__(
                 [
                     blocks.complex_to_interleaved_short(False, 32767),
                     blocks.stream_to_vector(gr.sizeof_short, fft_size * 2),
-                    iqtlabs.write_freq_samples(
+                    self.iqtlabs.write_freq_samples(
                         "rx_freq",
                         gr.sizeof_short,
                         "_".join(("ci16", endianstr())),
@@ -125,7 +128,7 @@ def __init__(
             logging.info(
                 f"retuning across {freq_range/1e6}MHz in {self.sweep_sec}s, requires retuning at {target_retune_hz}Hz in {tune_step_hz/1e6}MHz steps ({tune_step_fft} FFTs)"
             )
-        retune_fft = iqtlabs.retune_fft(
+        retune_fft = self.iqtlabs.retune_fft(
             "rx_freq",
             fft_size,
             fft_size,
@@ -135,7 +138,6 @@ def __init__(
             tune_step_hz,
             tune_step_fft,
             skip_tune_step,
-            fft_roll,
             db_clamp_floor,
             db_clamp_ceil,
             sample_dir,
@@ -165,7 +167,7 @@ def __init__(
             prediction_vlen = np.prod(prediction_shape)
             Path(inference_output_dir).mkdir(parents=True, exist_ok=True)
             Path(inference_output_dir, "images").mkdir(parents=True, exist_ok=True)
-            self.image_inference_block = iqtlabs.image_inference(
+            self.image_inference_block = self.iqtlabs.image_inference(
                 tag="rx_freq",
                 vlen=fft_size,
                 x=x,
@@ -265,30 +267,34 @@ def get_db_blocks(self, fft_size, samp_rate, scaling):
     def get_window(self, fft_size):
         return window.hann(fft_size)
 
-    def get_fft_blocks(self, fft_size, dc_block_len, dc_block_long):
+    def get_offload_fft_block(self, fft_batch_size, fft_size, fft_block):
+        return [
+            blocks.stream_to_vector(gr.sizeof_gr_complex, fft_batch_size * fft_size),
+            blocks.multiply_const_vff(
+                [val for val in self.get_window(fft_size) for _ in range(2)]
+                * fft_batch_size
+            ),
+            fft_block,
+            blocks.vector_to_stream(gr.sizeof_gr_complex * fft_size, fft_batch_size),
+            self.iqtlabs.vector_roll(fft_size),
+        ]
+
+    def get_fft_blocks(
+        self, use_vkfft, fft_batch_size, fft_size, dc_block_len, dc_block_long
+    ):
         fft_blocks = []
-        fft_roll = False
         if dc_block_len:
             fft_blocks.append(grfilter.dc_blocker_cc(dc_block_len, dc_block_long))
+        fft_block = None
         if self.wavelearner:
-            fft_batch_size = 256
-            fft_roll = True
+            fft_block = self.wavelearner.fft(
+                int(fft_batch_size * fft_size), (fft_size), True
+            )
+        elif use_vkfft:
+            fft_block = self.iqtlabs.vkfft(int(fft_batch_size * fft_size), fft_size)
+        if fft_block:
             fft_blocks.extend(
-                [
-                    blocks.stream_to_vector(
-                        gr.sizeof_gr_complex, fft_batch_size * fft_size
-                    ),
-                    blocks.multiply_const_vff(
-                        [val for val in self.get_window(fft_size) for _ in range(2)]
-                        * fft_batch_size
-                    ),
-                    self.wavelearner.fft(
-                        int(fft_batch_size * fft_size), (fft_size), True
-                    ),
-                    blocks.vector_to_stream(
-                        gr.sizeof_gr_complex * fft_size, fft_batch_size
-                    ),
-                ]
+                self.get_offload_fft_block(fft_batch_size, fft_size, fft_block)
             )
         else:
             fft_blocks.extend(
@@ -297,7 +303,7 @@ def get_fft_blocks(self, fft_size, dc_block_len, dc_block_long):
                     fft.fft_vcc(fft_size, True, self.get_window(fft_size), True, 1),
                 ]
             )
-        return (fft_blocks, fft_roll)
+        return fft_blocks
 
     def start(self):
         super().start()

diff --git a/gamutrf/scan.py b/gamutrf/scan.py
@@ -256,6 +256,20 @@ def argument_parser():
         action=BooleanOptionalAction,
         help="add sigmf meta file",
     )
+    parser.add_argument(
+        "--use_vkfft",
+        dest="use_vkfft",
+        default=True,
+        action=BooleanOptionalAction,
+        help="use VkFFT (ignored if wavelearner available)",
+    )
+    parser.add_argument(
+        "--fft_batch_size",
+        dest="fft_batch_size",
+        type=int,
+        default=256,
+        help="offload FFT batch size",
+    )
     return parser
 
 
@@ -331,6 +345,8 @@ def run_loop(options, prom_vars, wavelearner):
             scaling=handler.options.scaling,
             rotate_secs=handler.options.rotate_secs,
             description=handler.options.description,
+            fft_batch_size=handler.options.fft_batch_size,
+            use_vkfft=handler.options.use_vkfft,
             sigmf=handler.options.sigmf,
             iqtlabs=iqtlabs,
             wavelearner=wavelearner,