hip.cpp

Go to the documentation of this file.

/**
 * HIPGPUFFTGrid — GPUFFTGrid implementation using AMD hipFFT.
 *
 * Provides both FP64 (hipfftExecZ2Z) and FP32 (hipfftExecC2C) complex-to-complex
 * FFT transforms. hipFFT plans are created during init() and destroyed on
 * destruction.
 *
 * For the base FFTGrid interface (grid dimensions, g-vector mapping, buffers),
 * this class composes a SerialFFTGrid internally and delegates all base-class
 * pure virtual calls to it. This avoids diamond inheritance while reusing
 * the complete FFTW-based implementation from SerialFFTGrid.
 *
 * hipFFT operates on device pointers. With unified memory (hipMallocManaged),
 * the same pointers are accessible from both host and device.
 */
 
#include "layer_3_representation/basis/planewave/internal/fft/fft_grid_serial.hpp"
#include "layer_3_representation/basis/planewave/internal/fft/gpu/init.hpp"
 
#include <hip/hip_runtime.h>
#include <hipfft/hipfft.h>
#include <stdexcept>
#include <string>
 
namespace q3 {
namespace gpu {
 
namespace {
 
/**
 * Check a hipFFT return code and throw on error.
 */
void check_hipfft_error(hipfftResult result, const char* operation) {
    if (result != HIPFFT_SUCCESS) {
        throw std::runtime_error(std::string("HIPGPUFFTGrid: ") + operation +
                                 " failed with result " + std::to_string(static_cast<int>(result)));
    }
}
 
} // anonymous namespace
 
class HIPGPUFFTGrid : public GPUFFTGrid {
  public:
    /**
     * Construct a HIP FFT grid backed by the given device.
     *
     * hipFFT plans are not created here — they are created in init() after
     * the grid dimensions are known.
     */
    explicit HIPGPUFFTGrid(std::shared_ptr<execution::gpu::device::GPUDeviceBackend> device_backend)
        : GPUFFTGrid(std::move(device_backend)) {}
 
    ~HIPGPUFFTGrid() override {
        if (plan_fp64_ != 0) {
            hipfftDestroy(plan_fp64_);
        }
        if (plan_fp32_ != 0) {
            hipfftDestroy(plan_fp32_);
        }
        if (device_buffer_fp32_ != nullptr) {
            hipFree(device_buffer_fp32_);
        }
        if (device_buffer_fp64_ != nullptr) {
            hipFree(device_buffer_fp64_);
        }
        if (plan_fp32_batch_ != 0) {
            hipfftDestroy(plan_fp32_batch_);
        }
        if (device_buffer_fp32_batch_ != nullptr) {
            hipFree(device_buffer_fp32_batch_);
        }
        if (plan_fp64_batch_ != 0) {
            hipfftDestroy(plan_fp64_batch_);
        }
        if (device_buffer_fp64_batch_ != nullptr) {
            hipFree(device_buffer_fp64_batch_);
        }
    }
 
    // ── FFTGrid base interface (delegated to serial_fft_grid_) ───────────
 
    void init(int dimension_1,
              int dimension_2,
              int dimension_3,
              bool use_measured_planning = false) override {
        // Initialize the serial FFT grid (sets up dimensions and FFTW plans).
        serial_fft_grid_.init(dimension_1, dimension_2, dimension_3, use_measured_planning);
 
        // Create hipFFT FP64 plan (Z2Z = complex double-to-complex double).
        check_hipfft_error(
            hipfftPlan3d(&plan_fp64_, dimension_1, dimension_2, dimension_3, HIPFFT_Z2Z),
            "hipfftPlan3d FP64");
 
        // Create hipFFT FP32 plan (C2C = complex float-to-complex float).
        check_hipfft_error(
            hipfftPlan3d(&plan_fp32_, dimension_1, dimension_2, dimension_3, HIPFFT_C2C),
            "hipfftPlan3d FP32");
 
        // Allocate persistent device buffers for hipFFT.
        total_grid_points_cached_ = serial_fft_grid_.total_grid_points();
 
        const size_t buffer_size_fp32 =
            static_cast<size_t>(total_grid_points_cached_) * sizeof(hipfftComplex);
        hipMalloc(&device_buffer_fp32_, buffer_size_fp32);
 
        const size_t buffer_size_fp64 =
            static_cast<size_t>(total_grid_points_cached_) * sizeof(hipfftDoubleComplex);
        hipMalloc(&device_buffer_fp64_, buffer_size_fp64);
    }
 
    [[nodiscard]] int grid_dimension_1() const override {
        return serial_fft_grid_.grid_dimension_1();
    }
    [[nodiscard]] int grid_dimension_2() const override {
        return serial_fft_grid_.grid_dimension_2();
    }
    [[nodiscard]] int grid_dimension_3() const override {
        return serial_fft_grid_.grid_dimension_3();
    }
    [[nodiscard]] int total_grid_points() const override {
        return serial_fft_grid_.total_grid_points();
    }
    [[nodiscard]] int local_grid_points() const override {
        return serial_fft_grid_.local_grid_points();
    }
    [[nodiscard]] int local_offset() const override {
        return serial_fft_grid_.local_offset();
    }
 
    void forward(Complex* data) const override {
        serial_fft_grid_.forward(data);
    }
    void inverse(Complex* data) const override {
        serial_fft_grid_.inverse(data);
    }
    void inverse_raw(Complex* data) const override {
        serial_fft_grid_.inverse_raw(data);
    }
 
    [[nodiscard]] int
    g_vector_to_fft_grid_index(int miller_1, int miller_2, int miller_3) const override {
        return serial_fft_grid_.g_vector_to_fft_grid_index(miller_1, miller_2, miller_3);
    }
 
    void reduce_g_space(Complex* data, int number_of_elements) const override {
        serial_fft_grid_.reduce_g_space(data, number_of_elements);
    }
 
    Complex* input_buffer() const override {
        return serial_fft_grid_.input_buffer();
    }
    Complex* output_buffer() const override {
        return serial_fft_grid_.output_buffer();
    }
 
    void execute_inverse_raw_on_buffers() const override {
        serial_fft_grid_.execute_inverse_raw_on_buffers();
    }
    void execute_forward_on_buffers() const override {
        serial_fft_grid_.execute_forward_on_buffers();
    }
    void prepare_forward_from_output() const override {
        serial_fft_grid_.prepare_forward_from_output();
    }
 
    // ── FP32 FFT (hipFFT via persistent device buffer) ──────────────────────
 
    /**
     * Forward FP32 FFT via hipfftExecC2C with HIPFFT_FORWARD.
     *
     * Uses a persistent device buffer allocated in init() to avoid
     * per-call hipMalloc/hipFree overhead.
     */
    void forward_single(ComplexSingle* data) const override {
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * sizeof(hipfftComplex);
        hipMemcpy(device_buffer_fp32_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(
            hipfftExecC2C(plan_fp32_, device_buffer_fp32_, device_buffer_fp32_, HIPFFT_FORWARD),
            "hipfftExecC2C forward");
 
        hipMemcpy(data, device_buffer_fp32_, buffer_size, hipMemcpyDefault);
    }
 
    /**
     * Inverse FP32 FFT (raw, unnormalized) via hipfftExecC2C with HIPFFT_BACKWARD.
     *
     * Uses a persistent device buffer allocated in init() to avoid
     * per-call hipMalloc/hipFree overhead.
     */
    void inverse_raw_single(ComplexSingle* data) const override {
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * sizeof(hipfftComplex);
        hipMemcpy(device_buffer_fp32_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(
            hipfftExecC2C(plan_fp32_, device_buffer_fp32_, device_buffer_fp32_, HIPFFT_BACKWARD),
            "hipfftExecC2C inverse");
 
        hipMemcpy(data, device_buffer_fp32_, buffer_size, hipMemcpyDefault);
    }
 
    // ── FP64 FFT (hipFFT via persistent device buffer) ──────────────────────
 
    /**
     * Forward FP64 FFT via hipfftExecZ2Z with HIPFFT_FORWARD.
     *
     * Uses a persistent device buffer allocated in init() to avoid
     * per-call hipMalloc/hipFree overhead.
     */
    void forward_transform(Complex* data) const override {
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * sizeof(hipfftDoubleComplex);
        hipMemcpy(device_buffer_fp64_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(
            hipfftExecZ2Z(plan_fp64_, device_buffer_fp64_, device_buffer_fp64_, HIPFFT_FORWARD),
            "hipfftExecZ2Z forward");
 
        hipMemcpy(data, device_buffer_fp64_, buffer_size, hipMemcpyDefault);
    }
 
    /**
     * Inverse FP64 FFT (raw, unnormalized) via hipfftExecZ2Z with HIPFFT_BACKWARD.
     *
     * Uses a persistent device buffer allocated in init() to avoid
     * per-call hipMalloc/hipFree overhead.
     */
    void inverse_transform(Complex* data) const override {
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * sizeof(hipfftDoubleComplex);
        hipMemcpy(device_buffer_fp64_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(
            hipfftExecZ2Z(plan_fp64_, device_buffer_fp64_, device_buffer_fp64_, HIPFFT_BACKWARD),
            "hipfftExecZ2Z inverse");
 
        hipMemcpy(data, device_buffer_fp64_, buffer_size, hipMemcpyDefault);
    }
 
    /**
     * ── Batched FP32 FFT (hipfftPlanMany) ─────────────────────────────────
     *
     * Transforms multiple contiguous FFT grids in a single kernel launch.
     * The batch plan and device buffer are created lazily on first use and
     * resized if the batch size grows.
     */
 
    /**
     * Batched forward FP32 FFT via hipfftPlanMany + hipfftExecC2C.
     *
     * All batch_size grids are transformed in a single hipFFT kernel launch,
     * eliminating per-band dispatch overhead.
     */
    void forward_single_batch(ComplexSingle* data, int batch_size) const override {
        ensure_batch_resources(batch_size);
 
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * batch_size * sizeof(hipfftComplex);
        hipMemcpy(device_buffer_fp32_batch_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(hipfftExecC2C(plan_fp32_batch_,
                                         device_buffer_fp32_batch_,
                                         device_buffer_fp32_batch_,
                                         HIPFFT_FORWARD),
                           "hipfftExecC2C batched forward");
 
        hipMemcpy(data, device_buffer_fp32_batch_, buffer_size, hipMemcpyDefault);
    }
 
    /**
     * Batched inverse FP32 FFT (raw, unnormalized) via hipfftPlanMany + hipfftExecC2C.
     *
     * All batch_size grids are transformed in a single hipFFT kernel launch,
     * eliminating per-band dispatch overhead.
     */
    void inverse_raw_single_batch(ComplexSingle* data, int batch_size) const override {
        ensure_batch_resources(batch_size);
 
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * batch_size * sizeof(hipfftComplex);
        hipMemcpy(device_buffer_fp32_batch_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(hipfftExecC2C(plan_fp32_batch_,
                                         device_buffer_fp32_batch_,
                                         device_buffer_fp32_batch_,
                                         HIPFFT_BACKWARD),
                           "hipfftExecC2C batched inverse");
 
        hipMemcpy(data, device_buffer_fp32_batch_, buffer_size, hipMemcpyDefault);
    }
 
    // ── Batched FP64 FFT (hipfftPlanMany Z2Z) ─────────────────────────────
 
    /**
     * Batched forward FP64 FFT via hipfftPlanMany + hipfftExecZ2Z.
     *
     * All batch_size grids are transformed in a single hipFFT kernel launch.
     */
    void forward_transform_batch(Complex* data, int batch_size) const override {
        ensure_batch_resources_fp64(batch_size);
 
        const size_t buffer_size = static_cast<size_t>(total_grid_points_cached_) * batch_size *
                                   sizeof(hipfftDoubleComplex);
        hipMemcpy(device_buffer_fp64_batch_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(hipfftExecZ2Z(plan_fp64_batch_,
                                         device_buffer_fp64_batch_,
                                         device_buffer_fp64_batch_,
                                         HIPFFT_FORWARD),
                           "hipfftExecZ2Z batched forward");
 
        hipMemcpy(data, device_buffer_fp64_batch_, buffer_size, hipMemcpyDefault);
    }
 
    /**
     * Batched inverse FP64 FFT (raw, unnormalized) via hipfftPlanMany + hipfftExecZ2Z.
     *
     * All batch_size grids are transformed in a single hipFFT kernel launch.
     */
    void inverse_transform_batch(Complex* data, int batch_size) const override {
        ensure_batch_resources_fp64(batch_size);
 
        const size_t buffer_size = static_cast<size_t>(total_grid_points_cached_) * batch_size *
                                   sizeof(hipfftDoubleComplex);
        hipMemcpy(device_buffer_fp64_batch_, data, buffer_size, hipMemcpyDefault);
 
        check_hipfft_error(hipfftExecZ2Z(plan_fp64_batch_,
                                         device_buffer_fp64_batch_,
                                         device_buffer_fp64_batch_,
                                         HIPFFT_BACKWARD),
                           "hipfftExecZ2Z batched inverse");
 
        hipMemcpy(data, device_buffer_fp64_batch_, buffer_size, hipMemcpyDefault);
    }
 
  private:
    SerialFFTGrid serial_fft_grid_;
    hipfftHandle plan_fp64_ = 0;
    hipfftHandle plan_fp32_ = 0;
 
    // Persistent device buffers for single-grid hipFFT — allocated once in init().
    hipfftComplex* device_buffer_fp32_ = nullptr;
    hipfftDoubleComplex* device_buffer_fp64_ = nullptr;
    int total_grid_points_cached_ = 0;
 
    // Batched FP32 hipFFT resources — allocated lazily on first batch call.
    mutable hipfftHandle plan_fp32_batch_ = 0;
    mutable hipfftComplex* device_buffer_fp32_batch_ = nullptr;
    mutable int current_fp32_batch_size_ = 0;
 
    // Batched FP64 hipFFT resources — allocated lazily on first batch call.
    mutable hipfftHandle plan_fp64_batch_ = 0;
    mutable hipfftDoubleComplex* device_buffer_fp64_batch_ = nullptr;
    mutable int current_fp64_batch_size_ = 0;
 
    /**
     * Ensure the batched FP32 hipFFT plan and device buffer are large enough.
     */
    void ensure_batch_resources(int batch_size) const {
        if (batch_size == current_fp32_batch_size_) {
            return;
        }
 
        if (plan_fp32_batch_ != 0) {
            hipfftDestroy(plan_fp32_batch_);
            plan_fp32_batch_ = 0;
        }
        if (device_buffer_fp32_batch_ != nullptr) {
            hipFree(device_buffer_fp32_batch_);
            device_buffer_fp32_batch_ = nullptr;
        }
 
        int dimensions[3] = {serial_fft_grid_.grid_dimension_1(),
                             serial_fft_grid_.grid_dimension_2(),
                             serial_fft_grid_.grid_dimension_3()};
        int distance = total_grid_points_cached_;
 
        check_hipfft_error(hipfftPlanMany(&plan_fp32_batch_,
                                          3,           // rank (3D)
                                          dimensions,  // [d1, d2, d3]
                                          nullptr,     // inembed (null = contiguous)
                                          1,           // istride
                                          distance,    // idist
                                          nullptr,     // onembed (null = contiguous)
                                          1,           // ostride
                                          distance,    // odist
                                          HIPFFT_C2C,  // type
                                          batch_size), // batch
                           "hipfftPlanMany FP32 batch");
 
        const size_t buffer_size =
            static_cast<size_t>(total_grid_points_cached_) * batch_size * sizeof(hipfftComplex);
        hipMalloc(&device_buffer_fp32_batch_, buffer_size);
 
        current_fp32_batch_size_ = batch_size;
    }
 
    /**
     * Ensure the batched FP64 hipFFT plan and device buffer are large enough.
     */
    void ensure_batch_resources_fp64(int batch_size) const {
        if (batch_size == current_fp64_batch_size_) {
            return;
        }
 
        if (plan_fp64_batch_ != 0) {
            hipfftDestroy(plan_fp64_batch_);
            plan_fp64_batch_ = 0;
        }
        if (device_buffer_fp64_batch_ != nullptr) {
            hipFree(device_buffer_fp64_batch_);
            device_buffer_fp64_batch_ = nullptr;
        }
 
        int dimensions[3] = {serial_fft_grid_.grid_dimension_1(),
                             serial_fft_grid_.grid_dimension_2(),
                             serial_fft_grid_.grid_dimension_3()};
        int distance = total_grid_points_cached_;
 
        check_hipfft_error(hipfftPlanMany(&plan_fp64_batch_,
                                          3,           // rank (3D)
                                          dimensions,  // [d1, d2, d3]
                                          nullptr,     // inembed (null = contiguous)
                                          1,           // istride
                                          distance,    // idist
                                          nullptr,     // onembed (null = contiguous)
                                          1,           // ostride
                                          distance,    // odist
                                          HIPFFT_Z2Z,  // type
                                          batch_size), // batch
                           "hipfftPlanMany FP64 batch");
 
        const size_t buffer_size = static_cast<size_t>(total_grid_points_cached_) * batch_size *
                                   sizeof(hipfftDoubleComplex);
        hipMalloc(&device_buffer_fp64_batch_, buffer_size);
 
        current_fp64_batch_size_ = batch_size;
    }
};
 
std::unique_ptr<GPUFFTGrid>
create_hip_gpu_fft_grid(std::shared_ptr<execution::gpu::device::GPUDeviceBackend> device_backend) {
    return std::make_unique<HIPGPUFFTGrid>(std::move(device_backend));
}
 
} // namespace gpu
} // namespace q3