remvove callback params workspace

ASKabalan · ASKabalan · commit 9775bbab9dff · 2025-07-02T18:44:03.000+02:00
diff --git a/lib/src/extensions.cc b/lib/src/extensions.cc
@@ -65,14 +65,12 @@ constexpr bool is_double_v = is_double<T>::value;
  * @param input Input buffer containing HEALPix pixel-space data.
  * @param output Output buffer to store the FTM result.
  * @param workspace Output buffer for temporary workspace memory.
- * @param callback_params Output buffer for callback parameters.
  * @param descriptor Descriptor containing transform parameters.
  * @return ffi::Error indicating success or failure.
  */
 template <ffi::DataType T>
 ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Result<ffi::Buffer<T>> output,
                            ffi::Result<ffi::Buffer<T>> workspace,
-                           ffi::Result<ffi::Buffer<ffi::DataType::S64>> callback_params,
                            s2fftDescriptor descriptor) {
     // Step 1: Determine the complex type based on the XLA data type.
     using fft_complex_type = fft_complex_t<T>;
@@ -82,10 +80,9 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resul
     if (dim_in.size() == 2) {
         // Step 2a: Batched case.
         int batch_count = dim_in[0];
-        // Step 2b: Compute offsets for input, output, and callback parameters for each batch.
+        // Step 2b: Compute offsets for input and output for each batch.
         int64_t input_offset = descriptor.nside * descriptor.nside * 12;
         int64_t output_offset = (4 * descriptor.nside - 1) * (2 * descriptor.harmonic_band_limit);
-        int64_t params_offset = 2 * (descriptor.nside - 1) + 1;
 
         // Step 2c: Fork CUDA streams for parallel processing of batches.
         CudaStreamHandler handler;
@@ -99,16 +96,13 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resul
             auto executor = std::make_shared<s2fftExec<fft_complex_type>>();
             PlanCache::GetInstance().GetS2FFTExec(descriptor, executor);
 
-            // Step 2f: Calculate device pointers for the current batch's data, output, workspace, and
-            // callback parameters.
+            // Step 2f: Calculate device pointers for the current batch's data, output, and workspace.
             fft_complex_type* data_c =
                     reinterpret_cast<fft_complex_type*>(input.typed_data() + i * input_offset);
             fft_complex_type* out_c =
                     reinterpret_cast<fft_complex_type*>(output->typed_data() + i * output_offset);
             fft_complex_type* workspace_c =
                     reinterpret_cast<fft_complex_type*>(workspace->typed_data() + i * executor->m_work_size);
-            int64* callback_params_c =
-                    reinterpret_cast<int64*>(callback_params->typed_data() + i * params_offset);
 
             // Step 2g: Launch the forward transform on this sub-stream.
             executor->Forward(descriptor, sub_stream, data_c, workspace_c);
@@ -121,11 +115,10 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resul
         return ffi::Error::Success();
     } else {
         // Step 2j: Non-batched case.
-        // Step 2k: Get device pointers for data, output, workspace, and callback parameters.
+        // Step 2k: Get device pointers for data, output, and workspace.
         fft_complex_type* data_c = reinterpret_cast<fft_complex_type*>(input.typed_data());
         fft_complex_type* out_c = reinterpret_cast<fft_complex_type*>(output->typed_data());
         fft_complex_type* workspace_c = reinterpret_cast<fft_complex_type*>(workspace->typed_data());
-        int64* callback_params_c = reinterpret_cast<int64*>(callback_params->typed_data());
 
         // Step 2l: Get or create an s2fftExec instance from the PlanCache.
         auto executor = std::make_shared<s2fftExec<fft_complex_type>>();
@@ -152,14 +145,12 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resul
  * @param input Input buffer containing FTM data.
  * @param output Output buffer to store HEALPix pixel-space data.
  * @param workspace Output buffer for temporary workspace memory.
- * @param callback_params Output buffer for callback parameters.
  * @param descriptor Descriptor containing transform parameters.
  * @return ffi::Error indicating success or failure.
  */
 template <ffi::DataType T>
 ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Result<ffi::Buffer<T>> output,
                             ffi::Result<ffi::Buffer<T>> workspace,
-                            ffi::Result<ffi::Buffer<ffi::DataType::S64>> callback_params,
                             s2fftDescriptor descriptor) {
     // Step 1: Determine the complex type based on the XLA data type.
     using fft_complex_type = fft_complex_t<T>;
@@ -189,16 +180,13 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resu
             auto executor = std::make_shared<s2fftExec<fft_complex_type>>();
             PlanCache::GetInstance().GetS2FFTExec(descriptor, executor);
 
-            // Step 2f: Calculate device pointers for the current batch's data, output, workspace, and
-            // callback parameters.
+            // Step 2f: Calculate device pointers for the current batch's data, output, and workspace.
             fft_complex_type* data_c =
                     reinterpret_cast<fft_complex_type*>(input.typed_data() + i * input_offset);
             fft_complex_type* out_c =
                     reinterpret_cast<fft_complex_type*>(output->typed_data() + i * output_offset);
             fft_complex_type* workspace_c =
                     reinterpret_cast<fft_complex_type*>(workspace->typed_data() + i * executor->m_work_size);
-            int64* callback_params_c =
-                    reinterpret_cast<int64*>(callback_params->typed_data() + i * sizeof(int64) * 2);
 
             // Step 2g: Launch spectral folding kernel.
             s2fftKernels::launch_spectral_folding(data_c, out_c, descriptor.nside,
@@ -215,11 +203,10 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resu
         // Assertions to ensure correct input/output dimensions for non-batched operations.
         assert(dim_in.size() == 2);
         assert(dim_out.size() == 1);
-        // Step 2k: Get device pointers for data, output, workspace, and callback parameters.
+        // Step 2k: Get device pointers for data, output, and workspace.
         fft_complex_type* data_c = reinterpret_cast<fft_complex_type*>(input.typed_data());
         fft_complex_type* out_c = reinterpret_cast<fft_complex_type*>(output->typed_data());
         fft_complex_type* workspace_c = reinterpret_cast<fft_complex_type*>(workspace->typed_data());
-        int64* callback_params_c = reinterpret_cast<int64*>(callback_params->typed_data());
 
         // Step 2l: Get or create an s2fftExec instance from the PlanCache.
         auto executor = std::make_shared<s2fftExec<fft_complex_type>>();
@@ -310,24 +297,22 @@ s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, boo
  * @param input Input buffer.
  * @param output Output buffer.
  * @param workspace Output buffer for temporary workspace memory.
- * @param callback_params Output buffer for callback parameters.
  * @return ffi::Error indicating success or failure.
  */
 template <ffi::DataType T>
 ffi::Error healpix_fft_cuda(cudaStream_t stream, int64_t nside, int64_t harmonic_band_limit, bool reality,
                             bool forward, bool normalize, bool adjoint, ffi::Buffer<T> input,
-                            ffi::Result<ffi::Buffer<T>> output, ffi::Result<ffi::Buffer<T>> workspace,
-                            ffi::Result<ffi::Buffer<ffi::DataType::S64>> callback_params) {
+                            ffi::Result<ffi::Buffer<T>> output, ffi::Result<ffi::Buffer<T>> workspace) {
     // Step 1: Build the s2fftDescriptor based on the input parameters.
     size_t work_size = 0;  // Variable to hold the workspace size
     s2fftDescriptor descriptor = build_descriptor<T>(nside, harmonic_band_limit, reality, forward, normalize,
                                                      adjoint, true, work_size);
 
     // Step 2: Dispatch to either forward or backward transform based on the 'forward' flag.
     if (forward) {
-        return healpix_forward<T>(stream, input, output, workspace, callback_params, descriptor);
+        return healpix_forward<T>(stream, input, output, workspace, descriptor);
     } else {
-        return healpix_backward<T>(stream, input, output, workspace, callback_params, descriptor);
+        return healpix_backward<T>(stream, input, output, workspace, descriptor);
     }
 }
 
@@ -348,8 +333,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(healpix_fft_cuda_C64, healpix_fft_cuda<ffi::DataTy
                                       .Attr<bool>("adjoint")
                                       .Arg<ffi::Buffer<ffi::DataType::C64>>()
                                       .Ret<ffi::Buffer<ffi::DataType::C64>>()
-                                      .Ret<ffi::Buffer<ffi::DataType::C64>>()
-                                      .Ret<ffi::Buffer<ffi::DataType::S64>>());
+                                      .Ret<ffi::Buffer<ffi::DataType::C64>>());
 
 XLA_FFI_DEFINE_HANDLER_SYMBOL(healpix_fft_cuda_C128, healpix_fft_cuda<ffi::DataType::C128>,
                               ffi::Ffi::Bind()
@@ -362,8 +346,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(healpix_fft_cuda_C128, healpix_fft_cuda<ffi::DataT
                                       .Attr<bool>("adjoint")
                                       .Arg<ffi::Buffer<ffi::DataType::C128>>()
                                       .Ret<ffi::Buffer<ffi::DataType::C128>>()
-                                      .Ret<ffi::Buffer<ffi::DataType::C128>>()
-                                      .Ret<ffi::Buffer<ffi::DataType::S64>>());
+                                      .Ret<ffi::Buffer<ffi::DataType::C128>>());
 
 /**
  * @brief Encapsulates an FFI handler into a nanobind capsule.
diff --git a/s2fft/utils/healpix_ffts.py b/s2fft/utils/healpix_ffts.py
@@ -609,11 +609,7 @@ def _healpix_fft_cuda_abstract(f, L, nside, reality, fft_type, norm, adjoint):
         worksize //= 8  # 8 bytes per C64 element
         workspace_shape = (worksize,)
         workspace_dtype = np.complex64
-    # Step 3: Calculate shape for callback parameters.
-    nb_params = 2 * (nside - 1) + 1
-    params_shape = (nb_params,)
-
-    # Step 4: Define output shapes based on FFT type.
+    # Step 3: Define output shapes based on FFT type.
     healpix_size = (nside**2 * 12,)
     ftm_size = (4 * nside - 1, 2 * L)
     if fft_type == "forward":
@@ -628,17 +624,15 @@ def _healpix_fft_cuda_abstract(f, L, nside, reality, fft_type, norm, adjoint):
     else:
         raise ValueError(f"fft_type {fft_type} not recognised.")
 
-    # Step 5: Create ShapedArray objects for output, workspace, and callback parameters.
+    # Step 4: Create ShapedArray objects for output and workspace.
     workspace_aval = ShapedArray(
         shape=batch_shape + workspace_shape, dtype=workspace_dtype
     )
-    params_eval = ShapedArray(shape=batch_shape + params_shape, dtype=np.int64)
 
-    # Step 6: Return the ShapedArray objects.
+    # Step 5: Return the ShapedArray objects.
     return (
         f.update(shape=out_shape, dtype=f.dtype),
         workspace_aval,
-        params_eval,
     )
 
 
@@ -674,7 +668,7 @@ def _healpix_fft_cuda_lowering(ctx, f, *, L, nside, reality, fft_type, norm, adj
         raise MissingCUDASupport()
 
     # Step 2: Get the abstract evaluation results for the outputs.
-    (aval_out, _, _) = ctx.avals_out
+    (aval_out, _) = ctx.avals_out
 
     # Step 3: Get lowering information (double precision, forward/backward, normalize).
     is_double, forward, normalize = _get_lowering_info(fft_type, norm, aval_out.dtype)
@@ -839,8 +833,8 @@ def healpix_fft_cuda(
     """
     # Step 1: Promote input data to complex dtype if necessary.
     (f,) = promote_dtypes_complex(f)
-    # Step 2: Bind the input to the CUDA primitive. It returns multiple outputs (out, workspace, callback_params).
-    out, _, _ = _healpix_fft_cuda_primitive.bind(
+    # Step 2: Bind the input to the CUDA primitive. It returns multiple outputs (out, workspace).
+    out, _ = _healpix_fft_cuda_primitive.bind(
         f,
         L=L,
         nside=nside,
@@ -879,8 +873,8 @@ def healpix_ifft_cuda(
     """
     # Step 1: Promote input data to complex dtype if necessary.
     (ftm,) = promote_dtypes_complex(ftm)
-    # Step 2: Bind the input to the CUDA primitive. It returns multiple outputs (out, workspace, callback_params).
-    out, _, _ = _healpix_fft_cuda_primitive.bind(
+    # Step 2: Bind the input to the CUDA primitive. It returns multiple outputs (out, workspace).
+    out, _ = _healpix_fft_cuda_primitive.bind(
         ftm,
         L=L,
         nside=nside,
