ignis-numerics 0.0.1

data/lib/nvruby/fft/cufft_bindings.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+require "ffi"
+
+module Ignis
+  module FFT
+    # cuFFT library FFI bindings
+    module CuFFTBindings
+      extend FFI::Library
+
+      # FFT transform types
+      CUFFT_R2C = 0x2a  # Real to Complex (float)
+      CUFFT_C2R = 0x2c  # Complex to Real (float)
+      CUFFT_C2C = 0x29  # Complex to Complex (float)
+      CUFFT_D2Z = 0x6a  # Real to Complex (double)
+      CUFFT_Z2D = 0x6c  # Complex to Real (double)
+      CUFFT_Z2Z = 0x69  # Complex to Complex (double)
+
+      # FFT directions
+      CUFFT_FORWARD = -1
+      CUFFT_INVERSE = 1
+
+      # Callback types
+      CUFFT_CB_LD_COMPLEX = 0
+      CUFFT_CB_LD_COMPLEX_DOUBLE = 1
+      CUFFT_CB_LD_REAL = 2
+      CUFFT_CB_LD_REAL_DOUBLE = 3
+      CUFFT_CB_ST_COMPLEX = 4
+      CUFFT_CB_ST_COMPLEX_DOUBLE = 5
+      CUFFT_CB_ST_REAL = 6
+      CUFFT_CB_ST_REAL_DOUBLE = 7
+
+      @loaded = false
+
+      class << self
+        # Ensure cuFFT is loaded
+        # @return [void]
+        def ensure_loaded!
+          return if @loaded
+
+          CUDA::LibraryLoader.load_library(:cufft)
+
+          cuda_bin = Ignis.configuration.cuda_bin_path
+          if cuda_bin
+            ffi_lib Dir.glob(File.join(cuda_bin, "cufft64_*.dll")).max
+          else
+            ffi_lib "cufft64_11"
+          end
+
+          attach_cufft_functions!
+          @loaded = true
+        end
+
+        private
+
+        # rubocop:disable Metrics/MethodLength
+        def attach_cufft_functions!
+          # Plan management
+          attach_function :cufftCreate, [:pointer], :int
+          attach_function :cufftDestroy, [:pointer], :int
+          attach_function :cufftSetStream, [:pointer, :pointer], :int
+          attach_function :cufftGetSize, [:pointer, :pointer], :int
+
+          # 1D plan
+          attach_function :cufftPlan1d, [:pointer, :int, :int, :int], :int
+          # 2D plan
+          attach_function :cufftPlan2d, [:pointer, :int, :int, :int], :int
+          # 3D plan
+          attach_function :cufftPlan3d, [:pointer, :int, :int, :int, :int], :int
+          # Many (batched) plan
+          attach_function :cufftPlanMany, [
+            :pointer,      # plan
+            :int,          # rank
+            :pointer,      # n (dimensions)
+            :pointer,      # inembed
+            :int,          # istride
+            :int,          # idist
+            :pointer,      # onembed
+            :int,          # ostride
+            :int,          # odist
+            :int,          # type
+            :int           # batch
+          ], :int
+
+          # Plan with explicit memory
+          attach_function :cufftMakePlan1d, [:pointer, :int, :int, :int, :pointer], :int
+          attach_function :cufftMakePlan2d, [:pointer, :int, :int, :int, :pointer], :int
+          attach_function :cufftMakePlan3d, [:pointer, :int, :int, :int, :int, :pointer], :int
+          attach_function :cufftMakePlanMany, [
+            :pointer, :int, :pointer, :pointer, :int, :int,
+            :pointer, :int, :int, :int, :int, :pointer
+          ], :int
+
+          # Set workspace
+          attach_function :cufftSetAutoAllocation, [:pointer, :int], :int
+          attach_function :cufftSetWorkArea, [:pointer, :pointer], :int
+
+          # Execute - Complex to Complex (single precision)
+          attach_function :cufftExecC2C, [:pointer, :pointer, :pointer, :int], :int
+          # Execute - Complex to Complex (double precision)
+          attach_function :cufftExecZ2Z, [:pointer, :pointer, :pointer, :int], :int
+          # Execute - Real to Complex (single precision)
+          attach_function :cufftExecR2C, [:pointer, :pointer, :pointer], :int
+          # Execute - Complex to Real (single precision)
+          attach_function :cufftExecC2R, [:pointer, :pointer, :pointer], :int
+          # Execute - Real to Complex (double precision)
+          attach_function :cufftExecD2Z, [:pointer, :pointer, :pointer], :int
+          # Execute - Complex to Real (double precision)
+          attach_function :cufftExecZ2D, [:pointer, :pointer, :pointer], :int
+
+          # Estimate
+          attach_function :cufftEstimate1d, [:int, :int, :int, :pointer], :int
+          attach_function :cufftEstimate2d, [:int, :int, :int, :pointer], :int
+          attach_function :cufftEstimate3d, [:int, :int, :int, :int, :pointer], :int
+          attach_function :cufftEstimateMany, [
+            :int, :pointer, :pointer, :int, :int,
+            :pointer, :int, :int, :int, :int, :pointer
+          ], :int
+        end
+        # rubocop:enable Metrics/MethodLength
+      end
+
+      # Check cuFFT status and raise error if not success
+      # @param status [Integer] cuFFT status code
+      # @param context [String] Context for error message
+      # @return [void]
+      def self.check_status!(status, context = "cuFFT operation")
+        return if status.zero?
+
+        raise CuFFTError, status
+      end
+    end
+  end
+end

data/lib/nvruby/fft/fft_plan.rb

@@ -0,0 +1,288 @@
+# frozen_string_literal: true
+
+module Ignis
+  module FFT
+    # Reusable FFT plan for repeated operations with same dimensions
+    class FFTPlan
+      # @return [Array<Integer>] Input shape
+      attr_reader :shape
+
+      # @return [Symbol] Input data type
+      attr_reader :dtype
+
+      # @return [Symbol] Transform direction (:forward, :inverse)
+      attr_reader :direction
+
+      # @return [Symbol] Transform type (:c2c, :r2c, :c2r)
+      attr_reader :transform_type
+
+      # @return [FFI::Pointer] cuFFT plan handle
+      attr_reader :plan_handle
+
+      # @param shape [Array<Integer>] Input shape
+      # @param dtype [Symbol] Input data type
+      # @param direction [Symbol] :forward or :inverse
+      # @param transform_type [Symbol] :c2c, :r2c, or :c2r
+      # @param batch [Integer] Batch size for multiple FFTs
+      def initialize(shape:, dtype: :complex64, direction: :forward, transform_type: :c2c, batch: 1)
+        @shape = Array(shape)
+        @dtype = DType.validate!(dtype)
+        @direction = direction
+        @transform_type = transform_type
+        @batch = batch
+
+        @plan_handle = nil
+        @destroyed = false
+
+        create_plan!
+      end
+
+      # @return [Integer] Number of dimensions
+      def ndim
+        @shape.size
+      end
+
+      # @return [Array<Integer>] Output shape
+      def output_shape
+        case @transform_type
+        when :r2c
+          out = @shape.dup
+          out[-1] = out[-1] / 2 + 1
+          out
+        when :c2r
+          out = @shape.dup
+          # Output size should be provided, default to (n-1)*2
+          out[-1] = (@shape[-1] - 1) * 2
+          out
+        else
+          @shape.dup
+        end
+      end
+
+      # @return [Symbol] Output data type
+      def output_dtype
+        case @transform_type
+        when :r2c
+          DType.complex_dtype(@dtype)
+        when :c2r
+          DType.real_dtype(@dtype)
+        else
+          @dtype
+        end
+      end
+
+      # Execute the FFT plan
+      # @param input [NvArray] Input array
+      # @param output [NvArray, nil] Output array (created if nil)
+      # @param stream [CUDA::Stream, nil] CUDA stream
+      # @return [NvArray] Result
+      def execute(input, output: nil, stream: nil)
+        raise InvalidOperationError, "FFTPlan has been destroyed" if @destroyed
+
+        validate_input!(input)
+
+        input = input.to_device unless input.on_device?
+
+        output ||= NvArray.new(shape: output_shape, dtype: output_dtype, device: input.device_index)
+        output = output.to_device unless output.on_device?
+
+        if stream
+          status = CuFFTBindings.cufftSetStream(@plan_handle, stream.handle)
+          CuFFTBindings.check_status!(status, "Set FFT stream")
+        end
+
+        direction_flag = @direction == :forward ? CuFFTBindings::CUFFT_FORWARD : CuFFTBindings::CUFFT_INVERSE
+
+        status = case @transform_type
+                 when :c2c
+                   execute_c2c(input, output, direction_flag)
+                 when :r2c
+                   execute_r2c(input, output)
+                 when :c2r
+                   execute_c2r(input, output)
+                 else
+                   raise InvalidOperationError, "Unknown transform type: #{@transform_type}"
+                 end
+
+        CuFFTBindings.check_status!(status, "Execute FFT")
+
+        output
+      end
+
+      # Get estimated workspace size
+      # @return [Integer] Workspace size in bytes
+      def workspace_size
+        raise InvalidOperationError, "FFTPlan has been destroyed" if @destroyed
+
+        size_ptr = FFI::MemoryPointer.new(:size_t)
+        status = CuFFTBindings.cufftGetSize(@plan_handle, size_ptr)
+        CuFFTBindings.check_status!(status, "Get FFT workspace size")
+
+        size_ptr.read(:size_t)
+      end
+
+      # Destroy the plan and free resources
+      # @return [void]
+      def destroy!
+        return if @destroyed || @plan_handle.nil?
+
+        CuFFTBindings.cufftDestroy(@plan_handle)
+        @plan_handle = nil
+        @destroyed = true
+      end
+
+      # Check if plan has been destroyed
+      # @return [Boolean]
+      def destroyed?
+        @destroyed
+      end
+
+      # @return [String]
+      def to_s
+        status = @destroyed ? "destroyed" : "active"
+        "FFTPlan(shape=#{@shape}, dtype=#{@dtype}, #{@transform_type}, #{@direction}, #{status})"
+      end
+
+      private
+
+      # Create the cuFFT plan
+      def create_plan!
+        CuFFTBindings.ensure_loaded!
+
+        plan_ptr = FFI::MemoryPointer.new(:pointer)
+        fft_type = determine_fft_type
+
+        status = case ndim
+                 when 1
+                   CuFFTBindings.cufftPlan1d(plan_ptr, @shape[0], fft_type, @batch)
+                 when 2
+                   CuFFTBindings.cufftPlan2d(plan_ptr, @shape[0], @shape[1], fft_type)
+                 when 3
+                   CuFFTBindings.cufftPlan3d(plan_ptr, @shape[0], @shape[1], @shape[2], fft_type)
+                 else
+                   raise DimensionError, "FFT plan supports 1D, 2D, or 3D only, got #{ndim}D"
+                 end
+
+        CuFFTBindings.check_status!(status, "Create FFT plan")
+
+        @plan_handle = plan_ptr.read_pointer
+
+        ObjectSpace.define_finalizer(self, self.class.release_finalizer(@plan_handle))
+      end
+
+      # Determine cuFFT type constant
+      # @return [Integer]
+      def determine_fft_type
+        case @transform_type
+        when :c2c
+          case @dtype
+          when :complex64 then CuFFTBindings::CUFFT_C2C
+          when :complex128 then CuFFTBindings::CUFFT_Z2Z
+          else raise UnsupportedDTypeError.new(@dtype, operation: "FFT C2C")
+          end
+        when :r2c
+          case @dtype
+          when :float32 then CuFFTBindings::CUFFT_R2C
+          when :float64 then CuFFTBindings::CUFFT_D2Z
+          else raise UnsupportedDTypeError.new(@dtype, operation: "FFT R2C")
+          end
+        when :c2r
+          case @dtype
+          when :complex64 then CuFFTBindings::CUFFT_C2R
+          when :complex128 then CuFFTBindings::CUFFT_Z2D
+          else raise UnsupportedDTypeError.new(@dtype, operation: "FFT C2R")
+          end
+        else
+          raise InvalidOperationError, "Unknown transform type: #{@transform_type}"
+        end
+      end
+
+      # Validate input array
+      # For batched FFT, input can have shape [batch, *plan_shape] or [*plan_shape]
+      def validate_input!(input)
+        raise ArgumentError, "Expected NvArray, got #{input.class}" unless input.is_a?(NvArray)
+
+        # For batched execution, input can have extra leading dimensions
+        # The trailing dimensions must match the plan shape
+        plan_dims = @shape.size
+        input_dims = input.shape.size
+
+        if input_dims == plan_dims
+          # Direct match - single FFT or batch=1
+          unless input.shape == @shape
+            raise DimensionError, "Shape mismatch: expected #{@shape}, got #{input.shape}"
+          end
+        elsif input_dims == plan_dims + 1
+          # Batched input: [batch, *plan_shape]
+          trailing_shape = input.shape[-plan_dims..]
+          unless trailing_shape == @shape
+            raise DimensionError, "Shape mismatch: expected trailing dims #{@shape}, got #{trailing_shape}"
+          end
+          # Verify batch count matches plan batch
+          expected_batch = input.shape[0]
+          if @batch > 1 && expected_batch != @batch
+            raise DimensionError, "Batch mismatch: plan has #{@batch}, input has #{expected_batch}"
+          end
+        else
+          raise DimensionError, "Shape mismatch: expected #{plan_dims} or #{plan_dims + 1} dims, got #{input_dims}"
+        end
+
+        # Validate dtype
+        case @transform_type
+        when :c2c, :c2r
+          unless input.dtype == @dtype
+            raise ArgumentError, "Dtype mismatch: expected #{@dtype}, got #{input.dtype}"
+          end
+        when :r2c
+          real_dtype = DType.real_dtype(@dtype)
+          unless input.dtype == real_dtype || input.dtype == @dtype
+            raise ArgumentError, "Dtype mismatch: expected #{real_dtype}, got #{input.dtype}"
+          end
+        end
+      end
+
+      # Execute C2C transform
+      def execute_c2c(input, output, direction)
+        case @dtype
+        when :complex64
+          CuFFTBindings.cufftExecC2C(@plan_handle, input.device_ptr, output.device_ptr, direction)
+        when :complex128
+          CuFFTBindings.cufftExecZ2Z(@plan_handle, input.device_ptr, output.device_ptr, direction)
+        end
+      end
+
+      # Execute R2C transform
+      def execute_r2c(input, output)
+        real_dtype = DType.real_dtype(@dtype)
+        case real_dtype
+        when :float32
+          CuFFTBindings.cufftExecR2C(@plan_handle, input.device_ptr, output.device_ptr)
+        when :float64
+          CuFFTBindings.cufftExecD2Z(@plan_handle, input.device_ptr, output.device_ptr)
+        end
+      end
+
+      # Execute C2R transform
+      def execute_c2r(input, output)
+        case @dtype
+        when :complex64
+          CuFFTBindings.cufftExecC2R(@plan_handle, input.device_ptr, output.device_ptr)
+        when :complex128
+          CuFFTBindings.cufftExecZ2D(@plan_handle, input.device_ptr, output.device_ptr)
+        end
+      end
+
+      class << self
+        # Create finalizer for plan cleanup
+        # @param handle [FFI::Pointer] Plan handle
+        # @return [Proc]
+        def release_finalizer(handle)
+          proc do
+            CuFFTBindings.ensure_loaded!
+            CuFFTBindings.cufftDestroy(handle)
+          end
+        end
+      end
+    end
+  end
+end