ignis-numerics 0.0.1

data/lib/nvruby/fft/operations.rb

@@ -0,0 +1,364 @@
+# frozen_string_literal: true
+
+module Ignis
+  module FFT
+    # FFT operations using cuFFT
+    module Operations
+      class << self
+        # 1D Forward FFT
+        # @param x [NvArray] Input array
+        # @param axis [Integer] Axis along which to compute FFT (-1 for last)
+        # @param norm [Symbol] Normalization mode (:forward, :backward, :ortho)
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output
+        def fft(x, axis: -1, norm: :backward, stream: nil)
+          validate_input!(x)
+          axis = normalize_axis(axis, x.ndim)
+
+          n = x.shape[axis]
+          output_dtype = DType.complex_dtype(x.dtype)
+
+          result = execute_fft(x, output_dtype, :forward, [n], stream)
+          apply_normalization(result, n, :forward, norm)
+        end
+
+        # 1D Inverse FFT
+        # @param x [NvArray] Input complex array
+        # @param axis [Integer] Axis along which to compute IFFT
+        # @param norm [Symbol] Normalization mode (:forward, :backward, :ortho)
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output
+        def ifft(x, axis: -1, norm: :backward, stream: nil)
+          validate_complex_input!(x)
+          axis = normalize_axis(axis, x.ndim)
+
+          n = x.shape[axis]
+
+          result = execute_fft(x, x.dtype, :inverse, [n], stream)
+          apply_normalization(result, n, :inverse, norm)
+        end
+
+        # 2D Forward FFT
+        # @param x [NvArray] Input 2D array
+        # @param axes [Array<Integer>, nil] Axes for FFT (default last two)
+        # @param norm [Symbol] Normalization mode
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output
+        def fft2(x, axes: nil, norm: :backward, stream: nil)
+          validate_input!(x)
+          raise DimensionError, "Input must be at least 2D for fft2" if x.ndim < 2
+
+          axes ||= [-2, -1]
+          axes = axes.map { |a| normalize_axis(a, x.ndim) }
+
+          n = axes.map { |a| x.shape[a] }
+          output_dtype = DType.complex_dtype(x.dtype)
+
+          result = execute_fft(x, output_dtype, :forward, n, stream)
+          apply_normalization(result, n.reduce(:*), :forward, norm)
+        end
+
+        # 2D Inverse FFT
+        # @param x [NvArray] Input complex 2D array
+        # @param axes [Array<Integer>, nil] Axes for IFFT
+        # @param norm [Symbol] Normalization mode
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output
+        def ifft2(x, axes: nil, norm: :backward, stream: nil)
+          validate_complex_input!(x)
+          raise DimensionError, "Input must be at least 2D for ifft2" if x.ndim < 2
+
+          axes ||= [-2, -1]
+          axes = axes.map { |a| normalize_axis(a, x.ndim) }
+
+          n = axes.map { |a| x.shape[a] }
+
+          result = execute_fft(x, x.dtype, :inverse, n, stream)
+          apply_normalization(result, n.reduce(:*), :inverse, norm)
+        end
+
+        # N-dimensional FFT
+        # @param x [NvArray] Input array
+        # @param axes [Array<Integer>, nil] Axes for FFT
+        # @param norm [Symbol] Normalization mode
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output
+        def fftn(x, axes: nil, norm: :backward, stream: nil)
+          validate_input!(x)
+
+          axes ||= (0...x.ndim).to_a
+          axes = axes.map { |a| normalize_axis(a, x.ndim) }
+
+          n = axes.map { |a| x.shape[a] }
+          output_dtype = DType.complex_dtype(x.dtype)
+
+          result = execute_fft(x, output_dtype, :forward, n, stream)
+          apply_normalization(result, n.reduce(:*), :forward, norm)
+        end
+
+        # N-dimensional Inverse FFT
+        # @param x [NvArray] Input complex array
+        # @param axes [Array<Integer>, nil] Axes for IFFT
+        # @param norm [Symbol] Normalization mode
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output
+        def ifftn(x, axes: nil, norm: :backward, stream: nil)
+          validate_complex_input!(x)
+
+          axes ||= (0...x.ndim).to_a
+          axes = axes.map { |a| normalize_axis(a, x.ndim) }
+
+          n = axes.map { |a| x.shape[a] }
+
+          result = execute_fft(x, x.dtype, :inverse, n, stream)
+          apply_normalization(result, n.reduce(:*), :inverse, norm)
+        end
+
+        # Real-to-complex FFT
+        # @param x [NvArray] Real input array
+        # @param n [Integer, nil] FFT length (output length = n/2+1)
+        # @param axis [Integer] Axis for FFT
+        # @param norm [Symbol] Normalization mode
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Complex output (hermitian-symmetric)
+        def rfft(x, n: nil, axis: -1, norm: :backward, stream: nil)
+          validate_input!(x)
+          raise ArgumentError, "Input must be real for rfft" if DType.complex?(x.dtype)
+
+          axis = normalize_axis(axis, x.ndim)
+          n ||= x.shape[axis]
+
+          output_shape = x.shape.dup
+          output_shape[axis] = n / 2 + 1
+          output_dtype = DType.complex_dtype(x.dtype)
+
+          result = execute_rfft(x, output_shape, output_dtype, n, stream)
+          apply_normalization(result, n, :forward, norm)
+        end
+
+        # Complex-to-real Inverse FFT
+        # @param x [NvArray] Complex input array (hermitian-symmetric)
+        # @param n [Integer, nil] Output length
+        # @param axis [Integer] Axis for FFT
+        # @param norm [Symbol] Normalization mode
+        # @param stream [CUDA::Stream, nil] CUDA stream
+        # @return [NvArray] Real output
+        def irfft(x, n: nil, axis: -1, norm: :backward, stream: nil)
+          validate_complex_input!(x)
+
+          axis = normalize_axis(axis, x.ndim)
+          input_size = x.shape[axis]
+          n ||= (input_size - 1) * 2
+
+          output_shape = x.shape.dup
+          output_shape[axis] = n
+          output_dtype = DType.real_dtype(x.dtype)
+
+          result = execute_irfft(x, output_shape, output_dtype, n, stream)
+          apply_normalization(result, n, :inverse, norm)
+        end
+
+        private
+
+        # Validate input is an NvArray
+        def validate_input!(x)
+          raise ArgumentError, "Expected NvArray, got #{x.class}" unless x.is_a?(NvArray)
+          raise DimensionError, "Input must have at least 1 dimension" if x.ndim.zero?
+        end
+
+        # Validate input is complex
+        def validate_complex_input!(x)
+          validate_input!(x)
+          raise ArgumentError, "Expected complex dtype, got #{x.dtype}" unless DType.complex?(x.dtype)
+        end
+
+        # Normalize negative axis
+        # @param axis [Integer]
+        # @param ndim [Integer]
+        # @return [Integer]
+        def normalize_axis(axis, ndim)
+          axis += ndim if axis.negative?
+          raise DimensionError, "Axis #{axis} out of bounds for #{ndim}D array" unless axis >= 0 && axis < ndim
+
+          axis
+        end
+
+        # Execute FFT operation
+        # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+        def execute_fft(x, output_dtype, direction, dimensions, stream)
+          CuFFTBindings.ensure_loaded!
+
+          x = x.to_device unless x.on_device?
+
+          # Create output array
+          output = NvArray.new(shape: x.shape, dtype: output_dtype, device: x.device_index)
+          output.to_device
+
+          # Create plan based on dtype and dimensions
+          plan_ptr = FFI::MemoryPointer.new(:pointer)
+
+          case dimensions.size
+          when 1
+            fft_type = determine_fft_type(x.dtype, output_dtype)
+            status = CuFFTBindings.cufftPlan1d(plan_ptr, dimensions[0], fft_type, 1)
+          when 2
+            fft_type = determine_fft_type(x.dtype, output_dtype)
+            status = CuFFTBindings.cufftPlan2d(plan_ptr, dimensions[0], dimensions[1], fft_type)
+          when 3
+            fft_type = determine_fft_type(x.dtype, output_dtype)
+            status = CuFFTBindings.cufftPlan3d(plan_ptr, dimensions[0], dimensions[1], dimensions[2], fft_type)
+          else
+            raise NotImplementedError, "FFT for #{dimensions.size}D not yet implemented"
+          end
+
+          CuFFTBindings.check_status!(status, "Create FFT plan")
+
+          plan = plan_ptr.read_pointer
+
+          # Set stream if provided
+          if stream
+            status = CuFFTBindings.cufftSetStream(plan, stream.handle)
+            CuFFTBindings.check_status!(status, "Set FFT stream")
+          end
+
+          # Execute
+          dir = direction == :forward ? CuFFTBindings::CUFFT_FORWARD : CuFFTBindings::CUFFT_INVERSE
+
+          case output_dtype
+          when :complex64
+            status = CuFFTBindings.cufftExecC2C(plan, x.device_ffi_ptr, output.device_ffi_ptr, dir)
+          when :complex128
+            status = CuFFTBindings.cufftExecZ2Z(plan, x.device_ffi_ptr, output.device_ffi_ptr, dir)
+          else
+            raise UnsupportedDTypeError.new(output_dtype, operation: "FFT")
+          end
+
+          CuFFTBindings.check_status!(status, "Execute FFT")
+
+          # Cleanup
+          CuFFTBindings.cufftDestroy(plan)
+
+          output
+        end
+        # rubocop:enable Metrics/AbcSize, Metrics/MethodLength
+
+        # Execute real-to-complex FFT
+        def execute_rfft(x, output_shape, output_dtype, n, stream)
+          CuFFTBindings.ensure_loaded!
+
+          x = x.to_device unless x.on_device?
+
+          output = NvArray.new(shape: output_shape, dtype: output_dtype, device: x.device_index)
+          output.to_device
+
+          plan_ptr = FFI::MemoryPointer.new(:pointer)
+          fft_type = x.dtype == :float32 ? CuFFTBindings::CUFFT_R2C : CuFFTBindings::CUFFT_D2Z
+
+          status = CuFFTBindings.cufftPlan1d(plan_ptr, n, fft_type, 1)
+          CuFFTBindings.check_status!(status, "Create RFFT plan")
+
+          plan = plan_ptr.read_pointer
+
+          if stream
+            CuFFTBindings.cufftSetStream(plan, stream.handle)
+          end
+
+          status = if x.dtype == :float32
+                     CuFFTBindings.cufftExecR2C(plan, x.device_ffi_ptr, output.device_ffi_ptr)
+                   else
+                     CuFFTBindings.cufftExecD2Z(plan, x.device_ffi_ptr, output.device_ffi_ptr)
+                   end
+
+          CuFFTBindings.check_status!(status, "Execute RFFT")
+          CuFFTBindings.cufftDestroy(plan)
+
+          output
+        end
+
+        # Execute complex-to-real IFFT
+        def execute_irfft(x, output_shape, output_dtype, n, stream)
+          CuFFTBindings.ensure_loaded!
+
+          x = x.to_device unless x.on_device?
+
+          output = NvArray.new(shape: output_shape, dtype: output_dtype, device: x.device_index)
+          output.to_device
+
+          plan_ptr = FFI::MemoryPointer.new(:pointer)
+          fft_type = output_dtype == :float32 ? CuFFTBindings::CUFFT_C2R : CuFFTBindings::CUFFT_Z2D
+
+          status = CuFFTBindings.cufftPlan1d(plan_ptr, n, fft_type, 1)
+          CuFFTBindings.check_status!(status, "Create IRFFT plan")
+
+          plan = plan_ptr.read_pointer
+
+          if stream
+            CuFFTBindings.cufftSetStream(plan, stream.handle)
+          end
+
+          status = if output_dtype == :float32
+                     CuFFTBindings.cufftExecC2R(plan, x.device_ffi_ptr, output.device_ffi_ptr)
+                   else
+                     CuFFTBindings.cufftExecZ2D(plan, x.device_ffi_ptr, output.device_ffi_ptr)
+                   end
+
+          CuFFTBindings.check_status!(status, "Execute IRFFT")
+          CuFFTBindings.cufftDestroy(plan)
+
+          output
+        end
+
+        # Determine FFT type based on input/output dtypes
+        def determine_fft_type(input_dtype, output_dtype)
+          case [input_dtype, output_dtype]
+          when %i[float32 complex64], %i[complex64 complex64]
+            CuFFTBindings::CUFFT_C2C
+          when %i[float64 complex128], %i[complex128 complex128]
+            CuFFTBindings::CUFFT_Z2Z
+          else
+            raise UnsupportedDTypeError.new("#{input_dtype}->#{output_dtype}", operation: "FFT")
+          end
+        end
+
+        # Apply normalization to FFT result
+        # @param result [NvArray] FFT result
+        # @param n [Integer] Transform size
+        # @param direction [Symbol] :forward or :inverse
+        # @param norm [Symbol] :forward, :backward, or :ortho
+        # @return [NvArray] Normalized result
+        def apply_normalization(result, n, direction, norm)
+          scale = case norm
+                  when :backward
+                    direction == :inverse ? 1.0 / n : 1.0
+                  when :forward
+                    direction == :forward ? 1.0 / n : 1.0
+                  when :ortho
+                    1.0 / Math.sqrt(n)
+                  else
+                    1.0
+                  end
+
+          return result if scale == 1.0
+
+          # Actually apply the scale (previously a no-op, so ifft was off by N and
+          # ortho transforms off by sqrt(N); cuFFT does not normalize). Use cuBLAS
+          # scal over the raw float/double components (complex = 2 reals/element).
+          Ignis::LinAlg::CuBLASBindings.ensure_loaded!
+          handle = Ignis::LinAlg::CuBLASBindings.get_handle
+          count = DType.complex?(result.dtype) ? result.size * 2 : result.size
+          ptr = result.device_ffi_ptr
+
+          if %i[float64 complex128].include?(result.dtype)
+            alpha = FFI::MemoryPointer.new(:double).tap { |p| p.put_double(0, scale) }
+            status = Ignis::LinAlg::CuBLASBindings.cublasDscal_v2(handle, count, alpha, ptr, 1)
+          else
+            alpha = FFI::MemoryPointer.new(:float).tap { |p| p.put_float(0, scale) }
+            status = Ignis::LinAlg::CuBLASBindings.cublasSscal_v2(handle, count, alpha, ptr, 1)
+          end
+          Ignis::LinAlg::CuBLASBindings.check_status!(status, "FFT normalization scal")
+          result
+        end
+      end
+    end
+  end
+end

data/lib/nvruby/linalg/cutensor_bindings.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+require "ffi"
+
+module Ignis
+  module Linalg
+    # FFI bindings for NVIDIA cuTENSOR v2.4+
+    # High-performance tensor primitives for CUDA
+    module CuTensorBindings
+      extend FFI::Library
+
+      # cuTENSOR status codes
+      CUTENSOR_STATUS_SUCCESS = 0
+      CUTENSOR_STATUS_NOT_INITIALIZED = 1
+      CUTENSOR_STATUS_ALLOC_FAILED = 2
+      CUTENSOR_STATUS_INVALID_VALUE = 3
+      CUTENSOR_STATUS_ARCH_MISMATCH = 4
+      CUTENSOR_STATUS_MAPPING_ERROR = 5
+      CUTENSOR_STATUS_EXECUTION_FAILED = 6
+      CUTENSOR_STATUS_INTERNAL_ERROR = 7
+      CUTENSOR_STATUS_NOT_SUPPORTED = 8
+      CUTENSOR_STATUS_LICENSE_ERROR = 9
+      CUTENSOR_STATUS_CUBLAS_ERROR = 10
+      CUTENSOR_STATUS_CUDA_ERROR = 11
+      CUTENSOR_STATUS_INSUFFICIENT_WORKSPACE = 12
+      CUTENSOR_STATUS_INSUFFICIENT_DRIVER = 13
+
+      # cuTENSOR Algorithms
+      CUTENSOR_ALGO_DEFAULT = -1
+      CUTENSOR_ALGO_GETT = -2
+      CUTENSOR_ALGO_TGETT = -3
+
+      # Workspace preference enums
+      CUTENSOR_WORKSPACE_MIN = 1
+      CUTENSOR_WORKSPACE_RECOMMENDED = 2
+      CUTENSOR_WORKSPACE_MAX = 3
+
+      @loaded = false
+      @mutex = Mutex.new
+
+      class << self
+        # Ensure cuTENSOR library is loaded and functions attached
+        # @return [void]
+        def ensure_loaded!
+          @mutex.synchronize do
+            return if @loaded
+
+            CUDA::LibraryLoader.load_library(:cutensor)
+            dll_path = CUDA::LibraryLoader.library_paths[:cutensor]
+            
+            raise LibraryNotFoundError, "cutensor" unless dll_path
+
+            ffi_lib dll_path
+            attach_functions!
+
+            @loaded = true
+            Ignis.logger.debug("cuTENSOR bindings loaded from #{dll_path}")
+          end
+        end
+
+        # Check status and raise error if not success
+        # @param status [Integer] status code
+        # @param context [String] Context for error message
+        # @raise [CuTensorError] If status is not success
+        def check_status!(status, context = "cuTENSOR operation")
+          return if status == CUTENSOR_STATUS_SUCCESS
+
+          raise CuTensorError.new("#{context} failed", status_code: status)
+        end
+
+        private
+
+        def attach_functions!
+          # Handle management
+          attach_function :cutensorCreate, [:pointer], :int
+          attach_function :cutensorDestroy, [:pointer], :int
+
+          # Tensor descriptors
+          attach_function :cutensorCreateTensorDescriptor, [:pointer, :uint32, :pointer, :pointer, :int, :int], :int
+          attach_function :cutensorDestroyTensorDescriptor, [:pointer], :int
+
+          # Operation descriptors
+          attach_function :cutensorCreateContractionDescriptor, [:pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :int, :int], :int
+          attach_function :cutensorDestroyOperationDescriptor, [:pointer], :int
+
+          # Plan management
+          attach_function :cutensorCreatePlan, [:pointer, :pointer, :pointer, :int, :uint64], :int
+          attach_function :cutensorDestroyPlan, [:pointer], :int
+
+          # Execution
+          attach_function :cutensorContract, [:pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :uint64, :pointer], :int
+          attach_function :cutensorReduce, [:pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :pointer, :uint64, :pointer], :int
+        end
+      end
+    end
+
+    # Specialized error for cuTENSOR
+    class CuTensorError < StandardError
+      attr_reader :status_code
+
+      def initialize(message, status_code: nil)
+        @status_code = status_code
+        super("#{message} (Status: #{status_code})")
+      end
+    end
+  end
+end