convolver 1.0.1 → 3.0.1

data/ext/convolver/correlate_raw.c

@@ -0,0 +1,167 @@
+// ext/convolver/correlate_raw.c
+
+#include "correlate_raw.h"
+
+static inline size_t checked_add( size_t left, size_t right, const char *description ) {
+  if ( right > SIZE_MAX - left ) {
+    rb_raise( rb_eRangeError, "%s exceeds native implementation limit", description );
+  }
+  return left + right;
+}
+
+static inline size_t checked_multiply( size_t left, size_t right, const char *description ) {
+  if ( left != 0 && right > SIZE_MAX / left ) {
+    rb_raise( rb_eRangeError, "%s exceeds native implementation limit", description );
+  }
+  return left * right;
+}
+
+static inline size_t size_from_shape( int rank, const size_t *shape, const char *description ) {
+  size_t size = 1;
+  int i;
+  for ( i = 0; i < rank; i++ ) {
+    size = checked_multiply( size, shape[i], description );
+  }
+  return size;
+}
+
+// Sets reverse indices
+static inline void corner_reset( int rank, const size_t *shape, size_t *rev_indices ) {
+  int i;
+  for ( i = 0; i < rank; i++ ) { rev_indices[i] = shape[i] - 1; }
+}
+
+// Counts indices down, returns number of ranks that reset
+static inline int corner_dec( int rank, const size_t *shape, size_t *rev_indices ) {
+  int i = 0;
+  (void) rank;
+  while ( ! rev_indices[i]-- ) {
+    rev_indices[i] = shape[i] - 1;
+    i++;
+  }
+  return i;
+}
+
+// Generates co-increment steps by rank boundaries crossed, for the outer position as inner position is incremented by 1
+static inline void calc_co_increment(
+    int rank, const size_t *outer_shape, const size_t *inner_shape, size_t *co_increment ) {
+  size_t factor = 1;
+  int i;
+  co_increment[0] = 1; // co-increment is always 1 in lowest rank
+  for ( i = 0; i < rank; i++ ) {
+    size_t skipped = checked_multiply( factor, outer_shape[i] - inner_shape[i], "array offset" );
+    co_increment[i+1] = checked_add( co_increment[i], skipped, "array offset" );
+    factor = checked_multiply( factor, outer_shape[i], "array size" );
+  }
+}
+
+static inline size_t maximum_offset(
+    int rank, const size_t *outer_shape, const size_t *inner_shape, const char *description ) {
+  size_t factor = 1;
+  size_t offset = 0;
+  int i;
+  for ( i = 0; i < rank; i++ ) {
+    size_t dimension_offset = checked_multiply( factor, inner_shape[i] - 1, description );
+    offset = checked_add( offset, dimension_offset, description );
+    factor = checked_multiply( factor, outer_shape[i], "array size" );
+  }
+  return offset;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+//  Correlate
+//
+
+void correlate_raw(
+    int in_rank, const size_t *in_shape, const float *in_ptr,
+    int kernel_rank, const size_t *kernel_shape, const float *kernel_ptr,
+    int out_rank, const size_t *out_shape, float *out_ptr ) {
+  size_t i, j, input_size, kernel_size, kernel_aligned, out_size, offset;
+  size_t maximum_input_offset;
+  size_t out_co_incr[LARGEST_RANK + 1], kernel_co_incr[LARGEST_RANK + 1];
+  size_t ker_q[LARGEST_RANK], out_q[LARGEST_RANK];
+  size_t *kernel_co_incr_cache;
+  VALUE cache_storage = 0;
+
+  kernel_size = size_from_shape( kernel_rank, kernel_shape, "kernel size" );
+  kernel_aligned = kernel_size - kernel_size % 4;
+  out_size = size_from_shape( out_rank, out_shape, "output size" );
+  input_size = size_from_shape( in_rank, in_shape, "input size" );
+
+  calc_co_increment( in_rank, in_shape, out_shape, out_co_incr );
+  calc_co_increment( in_rank, in_shape, kernel_shape, kernel_co_incr );
+  maximum_input_offset = checked_add(
+    maximum_offset( in_rank, in_shape, out_shape, "input offset" ),
+    maximum_offset( in_rank, in_shape, kernel_shape, "input offset" ),
+    "input offset"
+  );
+  if ( maximum_input_offset >= input_size ) {
+    rb_raise( rb_eRangeError, "input offset exceeds native implementation limit" );
+  }
+
+  kernel_co_incr_cache = RB_ALLOCV_N( size_t, cache_storage, kernel_size );
+  kernel_co_incr_cache[0] = 0;
+
+  corner_reset( kernel_rank, kernel_shape, ker_q );
+  for ( i = 1; i < kernel_size; i++ ) {
+    kernel_co_incr_cache[i] = checked_add(
+      kernel_co_incr_cache[i-1],
+      kernel_co_incr[ corner_dec( kernel_rank, kernel_shape, ker_q ) ],
+      "kernel offset"
+    );
+  }
+
+  offset = 0;
+  corner_reset( out_rank, out_shape, out_q );
+
+  // Main correlation loop
+  for ( i = 0; i < out_size; i++ ) {
+#if CONVOLVER_USE_SSE
+    __m128 simd_x, simd_y, simd_t;
+    float v[4];
+    simd_t = _mm_setzero_ps();
+#endif
+    float t = 0.0;
+
+#if CONVOLVER_USE_SSE
+    // Contiguous kernel views need not start on a 16-byte boundary.
+    for ( j = 0; j < kernel_aligned; j +=4 ) {
+      simd_x = _mm_loadu_ps( kernel_ptr + j );
+      // Yes the backwards alignment is correct
+      simd_y = _mm_set_ps( in_ptr[ offset + kernel_co_incr_cache[j+3] ], in_ptr[ offset + kernel_co_incr_cache[j+2] ],
+                           in_ptr[ offset + kernel_co_incr_cache[j+1] ], in_ptr[ offset + kernel_co_incr_cache[j] ] );
+      simd_x = _mm_mul_ps( simd_x, simd_y );
+      simd_t = _mm_add_ps( simd_x, simd_t );
+    }
+    _mm_store_ps( v, simd_t );
+#else
+    // SSE is unavailable on architectures such as Apple Silicon.
+    // Use the same calculation without x86-specific intrinsics.
+    for ( j = 0; j < kernel_aligned; j++ ) {
+      t += in_ptr[ offset + kernel_co_incr_cache[j] ] * kernel_ptr[j];
+    }
+#endif
+
+    // Complete any remaining 1,2 or 3 items one at a time
+    for ( j = kernel_aligned; j < kernel_size; j++ ) {
+      t += in_ptr[ offset + kernel_co_incr_cache[j] ] * kernel_ptr[ j ];
+    }
+
+#if CONVOLVER_USE_SSE
+    out_ptr[i] = v[0] + v[1] + v[2] + v[3] + t;
+#else
+    out_ptr[i] = t;
+#endif
+
+    if ( i + 1 < out_size ) {
+      offset = checked_add(
+        offset,
+        out_co_incr[ corner_dec( out_rank, out_shape, out_q ) ],
+        "input offset"
+      );
+    }
+  }
+
+  RB_ALLOCV_END( cache_storage );
+}

data/ext/convolver/correlate_raw.h

@@ -0,0 +1,28 @@
+// ext/convolver/correlate_raw.h
+
+////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Declarations of narray helper functions
+//
+
+#ifndef CORRELATE_RAW_H
+#define CORRELATE_RAW_H
+
+#ifndef CONVOLVER_RAW_CONFIG
+#define CONVOLVER_RAW_CONFIG
+#include <ruby.h>
+#if defined(__SSE__)
+#include <xmmintrin.h>
+#define CONVOLVER_USE_SSE 1
+#else
+#define CONVOLVER_USE_SSE 0
+#endif
+#define LARGEST_RANK 16
+#endif
+
+void correlate_raw(
+    int in_rank, const size_t *in_shape, const float *in_ptr,
+    int kernel_rank, const size_t *kernel_shape, const float *kernel_ptr,
+    int out_rank, const size_t *out_shape, float *out_ptr );
+
+#endif

data/lib/convolver/buffer_limits.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Convolver
+  # Checks individual allocations against native size and tagged-stride limits.
+  # @private
+  class BufferLimits
+    SIZE_MAX = (1 << ([0].pack('J').bytesize * 8)) - 1
+    # Numo tags signed strides with a left shift: floor(SSIZE_MAX / 2).
+    BYTE_MAX = SIZE_MAX / 4
+    INT_MAX = (1 << (([0].pack('i').bytesize * 8) - 1)) - 1
+    INDEX_BYTES = [0].pack('J').bytesize
+
+    attr_reader :size_max, :byte_max, :int_max
+
+    def initialize(size_max: SIZE_MAX, byte_max: BYTE_MAX, int_max: INT_MAX)
+      @size_max = size_max
+      @byte_max = byte_max
+      @int_max = int_max
+    end
+
+    def buffer!(count, width, description, error: RangeError)
+      return if count <= byte_max / width
+
+      raise error, "#{description} exceeds native byte/stride limit"
+    end
+
+    def product!(shape, description, error: RangeError)
+      shape.reduce(1) do |product, size|
+        raise error, "#{description} exceeds native element limit" if size > size_max / product
+
+        product * size
+      end
+    end
+
+    NATIVE = new.freeze
+  end
+
+  private_constant :BufferLimits
+end

data/lib/convolver/buffer_preparation.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'convolver/fft_limits'
+
+module Convolver
+  # Preflights casts, extension, indexing and direct-path cache allocations.
+  # @private
+  class BufferPreparation
+    def initialize(plan, signal_shape, kernel_shape, limits: BufferLimits::NATIVE)
+      @plan = plan
+      @signal_size = limits.product!(signal_shape, 'signal size')
+      @kernel_size = limits.product!(kernel_shape, 'kernel size')
+      @limits = limits
+      @scalar = signal_shape.empty?
+    end
+
+    def basic!
+      common!
+      extension!
+      limits.buffer!(kernel_size, BufferLimits::INDEX_BYTES, 'direct kernel offset cache')
+    end
+
+    def fft!
+      common!
+      return if @scalar
+      return extension! unless plan.wrap?
+
+      limits.buffer!(signal_size, 8, 'circular signal preparation', error: FftUnavailable)
+      limits.buffer!(kernel_size, 8, 'circular kernel preparation', error: FftUnavailable)
+      limits.buffer!(kernel_size, BufferLimits::INDEX_BYTES, 'circular kernel index buffer', error: FftUnavailable)
+    end
+
+    private
+
+    attr_reader :plan, :signal_size, :kernel_size, :limits
+
+    def common!
+      limits.buffer!(plan.result_size, 4, 'result buffer')
+      limits.buffer!(signal_size, 4, 'signal conversion buffer')
+      limits.buffer!(kernel_size, 4, 'kernel conversion buffer')
+    end
+
+    def extension!
+      return if plan.valid?
+
+      limits.buffer!(plan.extended_size, 4, 'extended signal buffer')
+      return if plan.boundary == :constant
+
+      plan.extended_shape.each do |length|
+        limits.buffer!(length, BufferLimits::INDEX_BYTES, 'boundary index buffer')
+      end
+    end
+  end
+
+  private_constant :BufferPreparation
+end

data/lib/convolver/circular_fft_operation.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module Convolver
+  # Calculates same-sized periodic operations through PocketFFT transforms.
+  # @private
+  class CircularFftOperation
+    def initialize(operation, signal, kernel, plan)
+      @real_axis = plan.fft_buffers.real_axis
+      @operation = operation
+      @signal = Numo::DFloat.cast(signal)
+      @kernel = fold_kernel(Numo::DFloat.cast(kernel), signal.shape, plan.anchors)
+    end
+
+    def call
+      return complex_product unless real_axis
+
+      real_product(real_axis)
+    end
+
+    def self.real_axis(shape)
+      return nil if shape.empty?
+
+      final_axis = shape.length - 1
+      return final_axis if shape[final_axis].even?
+
+      shape.each_index.select { |axis| shape[axis].even? }.max_by { |axis| shape[axis] }
+    end
+
+    private
+
+    attr_reader :operation, :signal, :kernel, :real_axis
+
+    def real_product(real_axis)
+      final_axis = signal.ndim - 1
+      prepared_signal = prepare_real_operand(signal, real_axis, final_axis)
+      prepared_kernel = prepare_real_operand(kernel, real_axis, final_axis)
+      result = real_spectrum_product(prepared_signal, prepared_kernel)
+      result = result.swapaxes(real_axis, final_axis) unless real_axis == final_axis
+      Numo::SFloat.cast(result)
+    end
+
+    def prepare_real_operand(value, real_axis, final_axis)
+      real_axis == final_axis ? value : value.swapaxes(real_axis, final_axis)
+    end
+
+    def real_spectrum_product(signal_value, kernel_value)
+      signal_spectrum = Numo::Pocketfft.rfftn(signal_value)
+      kernel_spectrum = Numo::Pocketfft.rfftn(kernel_value)
+      kernel_spectrum.inplace.conj if operation == :correlation
+      Numo::Pocketfft.irfftn(signal_spectrum.inplace * kernel_spectrum)
+    end
+
+    def complex_product
+      signal_spectrum = Numo::Pocketfft.fftn(signal)
+      kernel_spectrum = Numo::Pocketfft.fftn(kernel)
+      kernel_spectrum.inplace.conj if operation == :correlation
+      result = Numo::Pocketfft.ifftn(signal_spectrum.inplace * kernel_spectrum).real
+      Numo::SFloat.cast(result)
+    end
+
+    def fold_kernel(source, signal_shape, anchors)
+      folded = Numo::DFloat.zeros(*signal_shape)
+      source.flatten.each_with_index do |value, flat_index|
+        target = folded_target(flat_index, source.shape, signal_shape, anchors)
+        folded[*target] = folded[*target] + value
+      end
+      folded
+    end
+
+    def folded_target(flat_index, kernel_shape, signal_shape, anchors)
+      remainder = flat_index
+      Array.new(kernel_shape.length).tap do |target|
+        (kernel_shape.length - 1).downto(0) do |axis|
+          coordinate = remainder % kernel_shape[axis]
+          remainder /= kernel_shape[axis]
+          target[axis] = (coordinate - anchors[axis]) % signal_shape[axis]
+        end
+      end
+    end
+  end
+
+  private_constant :CircularFftOperation
+end

data/lib/convolver/fft_buffer_plan.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require 'convolver/fft_limits'
+require 'convolver/buffer_preparation'
+
+module Convolver
+  # Shares allocation-free FFT dimensions between execution and estimation.
+  # @private
+  class FftBufferPlan
+    attr_reader :transform_shape, :real_axis, :transform_size, :spectrum_size
+
+    def initialize(plan, signal_shape, kernel_shape, limits: BufferLimits::NATIVE)
+      BufferPreparation.new(plan, signal_shape, kernel_shape, limits:).fft!
+      @real_axis = plan.wrap? ? CircularFftOperation.real_axis(signal_shape) : signal_shape.length - 1
+      @transform_shape = select_shape(plan, signal_shape, kernel_shape, FftLimits.new(limits))
+      @transform_size = transform_shape.reduce(1, :*)
+      @spectrum_size = spectrum_elements
+    end
+
+    private
+
+    def select_shape(plan, signal_shape, kernel_shape, limits)
+      return [].freeze if signal_shape.empty?
+      return plan.linear_fft_shape(kernel_shape, limits:) unless plan.wrap?
+
+      limits.validate!(signal_shape, real_axis:)
+      signal_shape.dup.freeze
+    end
+
+    def spectrum_elements
+      return transform_size if transform_shape.empty? || real_axis.nil?
+
+      length = transform_shape[real_axis]
+      transform_size / length * ((length / 2) + 1)
+    end
+  end
+
+  private_constant :FftBufferPlan
+end

data/lib/convolver/fft_estimator.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Convolver
+  # Estimates the selected real or complex PocketFFT implementation cost.
+  # @private
+  class FftEstimator
+    REAL_FFT_COST = { 1 => 1.0e-9, 2 => 9.0e-10, 3 => 1.45e-9 }.freeze
+    COMPLEX_FFT_COST = { 1 => 2.4e-9, 2 => 2.4e-9, 3 => 2.8e-9 }.freeze
+    FIXED_COST = 1.5e-4
+    PREPARATION_COST = 1.0e-9
+    AXIS_MOVE_COST = 1.6e-9
+    CORRELATION_COST = 2.0e-10
+
+    def initialize(operation, signal, kernel, plan)
+      @operation = operation
+      @signal = signal
+      @kernel = kernel
+      @plan = plan
+      @buffers = plan.fft_buffers
+    end
+
+    def call
+      transform_size, spectrum_size, coefficient, moved_size = dimensions
+      transform_cost = coefficient * transform_size * Math.log([transform_size, 2].max)
+      preparation_cost = PREPARATION_COST * preparation_size(transform_size)
+      axis_move_cost = AXIS_MOVE_COST * moved_size
+      correlation_cost = operation == :correlation ? CORRELATION_COST * spectrum_size : 0.0
+      FIXED_COST + transform_cost + preparation_cost + axis_move_cost + correlation_cost
+    end
+
+    private
+
+    attr_reader :operation, :signal, :kernel, :plan, :buffers
+
+    def dimensions
+      return linear_dimensions unless plan.wrap?
+
+      circular_dimensions
+    end
+
+    def circular_dimensions
+      real_axis = buffers.real_axis
+      [buffers.transform_size, buffers.spectrum_size, transform_coefficient(real_axis), moved_size(real_axis)]
+    end
+
+    def linear_dimensions
+      [buffers.transform_size, buffers.spectrum_size,
+       cost_for(REAL_FFT_COST), 0]
+    end
+
+    def preparation_size(transform_size)
+      extension_size = plan.valid? || plan.wrap? ? 0 : plan.extended_size
+      working_size = plan.wrap? ? signal.size + kernel.size : 2 * transform_size
+      extension_size + working_size
+    end
+
+    def cost_for(costs)
+      costs.fetch(signal.ndim, costs.values.last)
+    end
+
+    def transform_coefficient(real_axis)
+      cost_for(real_axis ? REAL_FFT_COST : COMPLEX_FFT_COST)
+    end
+
+    def moved_size(real_axis)
+      real_axis && real_axis != signal.ndim - 1 ? 3 * signal.size : 0
+    end
+  end
+
+  private_constant :FftEstimator
+end

data/lib/convolver/fft_limits.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require 'convolver/buffer_limits'
+
+# Internal allocation and arithmetic checks for the PocketFFT paths.
+module Convolver
+  # Only FFT representability failures permit automatic direct fallback.
+  # @private
+  class FftUnavailable < RangeError; end
+
+  # Bounds PocketFFT buffers and signed arithmetic without constructing plans.
+  # @private
+  class FftLimits
+    def initialize(buffers = BufferLimits::NATIVE)
+      @buffers = buffers
+    end
+
+    def validate!(shape, real_axis:)
+      total = buffers.product!(shape, 'FFT size', error: FftUnavailable)
+      validate_buffers!(shape, total, real_axis)
+      # The wrapper casts total to int before dividing it by the axis length.
+      integer!(total, buffers.int_max, 'FFT total size')
+      # Real forward computes int z_step = (z / 2 + 1) * 2.
+      integer!(shape[real_axis], buffers.int_max - 2, 'FFT real stride') if real_axis
+      validate_axes!(shape)
+    end
+
+    private
+
+    attr_reader :buffers
+
+    def validate_axes!(shape)
+      # FFTPACK twiddle counts are <= 2n complex / 3n real. Bluestein m < 4n,
+      # so each variable allocation is <= 128n bytes, including inner twiddles.
+      # The 16n int bound also protects signed intermediates at the inner m.
+      shape.each do |length|
+        integer!(length, buffers.int_max / 16, 'FFT axis arithmetic')
+        buffer!(length, 128, 'FFT axis work buffer')
+      end
+    end
+
+    def validate_buffers!(shape, total, real_axis)
+      buffer!(total, 8, 'FFT spatial buffer')
+      return buffer!(total, 16, 'FFT complex buffer') unless real_axis
+
+      spectrum = total / shape[real_axis] * ((shape[real_axis] / 2) + 1)
+      buffer!(spectrum, 16, 'FFT spectrum buffer')
+      buffer!(total, 16, 'FFT inverse staging buffer')
+      buffer!(total, 8, 'FFT inverse output buffer')
+    end
+
+    def buffer!(count, width, description)
+      buffers.buffer!(count, width, description, error: FftUnavailable)
+    end
+
+    def integer!(value, maximum, description)
+      return if value <= maximum
+
+      raise FftUnavailable, "#{description} exceeds native integer limit"
+    end
+  end
+
+  private_constant :FftUnavailable, :FftLimits
+end

data/lib/convolver/linear_fft_operation.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Convolver
+  # Calculates a valid linear operation through real PocketFFT transforms.
+  # @private
+  class LinearFftOperation
+    def initialize(operation, signal, kernel, plan)
+      @operation = operation
+      @signal = signal
+      @kernel = kernel
+      @transform_shape = plan.fft_buffers.transform_shape
+    end
+
+    def call
+      signal_padded = zero_padded(signal)
+      kernel_padded = zero_padded(kernel)
+      full_result = real_spectrum_product(signal_padded, kernel_padded)
+      Numo::SFloat.cast(full_result[*valid_ranges])
+    end
+
+    private
+
+    attr_reader :operation, :signal, :kernel, :transform_shape
+
+    def zero_padded(value)
+      Numo::DFloat.zeros(*transform_shape).tap do |padded|
+        padded[*value.shape.map { |size| 0...size }] = value
+      end
+    end
+
+    def real_spectrum_product(signal_value, kernel_value)
+      signal_spectrum = Numo::Pocketfft.rfftn(signal_value)
+      kernel_spectrum = Numo::Pocketfft.rfftn(kernel_value)
+      kernel_spectrum.inplace.conj if operation == :correlation
+      Numo::Pocketfft.irfftn(signal_spectrum.inplace * kernel_spectrum)
+    end
+
+    def valid_ranges
+      signal.shape.zip(kernel.shape).map do |signal_size, kernel_size|
+        if operation == :correlation
+          0...(signal_size - kernel_size + 1)
+        else
+          (kernel_size - 1)...signal_size
+        end
+      end
+    end
+  end
+
+  private_constant :LinearFftOperation
+end