diff --git a/qdp/qdp-core/src/gpu/encodings/iqp.rs b/qdp/qdp-core/src/gpu/encodings/iqp.rs index 33d18cfaf0..c15d801000 100644 --- a/qdp/qdp-core/src/gpu/encodings/iqp.rs +++ b/qdp/qdp-core/src/gpu/encodings/iqp.rs @@ -64,6 +64,114 @@ impl IqpEncoder { num_qubits } } + + /// Encode multiple IQP samples using Matrix-Free Implicit Hadamard Tensor Core engine + #[cfg(target_os = "linux")] + pub fn encode_batch_tc( + &self, + device: &Arc, + batch_data: &[f64], + num_samples: usize, + sample_size: usize, + num_qubits: usize, + ) -> Result { + validate_qubit_count(num_qubits)?; + let expected_len = self.expected_data_len(num_qubits); + + if sample_size != expected_len { + return Err(MahoutError::InvalidInput(format!( + "Batch sample size {} does not match required parameters {}", + sample_size, expected_len + ))); + } + + if batch_data.len() != num_samples * sample_size { + return Err(MahoutError::InvalidInput(format!( + "Batch data length {} does not match num_samples {} * sample_size {}", + batch_data.len(), + num_samples, + sample_size + ))); + } + + for (i, &val) in batch_data.iter().enumerate() { + if !val.is_finite() { + let sample_idx = i / sample_size; + let param_idx = i % sample_size; + return Err(MahoutError::InvalidInput(format!( + "Sample {} parameter {} must be finite, got {}", + sample_idx, param_idx, val + ))); + } + } + + let state_len = 1 << num_qubits; + + let batch_state_vector = { + crate::profile_scope!("GPU::AllocBatchTC"); + GpuStateVector::new_batch(device, num_samples, num_qubits, Precision::Float64)? + }; + + let input_bytes = std::mem::size_of_val(batch_data); + let data_gpu = { + crate::profile_scope!("GPU::H2D_BatchIqpDataTC"); + device.htod_sync_copy(batch_data).map_err(|e| { + map_allocation_error(input_bytes, "IQP TC batch upload", Some(num_qubits), e) + })? + }; + + let state_ptr = batch_state_vector.ptr_f64().ok_or_else(|| { + MahoutError::InvalidInput( + "Batch state vector precision mismatch (expected float64 buffer)".to_string(), + ) + })?; + + { + crate::profile_scope!("GPU::BatchKernelLaunchTC"); + let ret = unsafe { + qdp_kernels::launch_iqp_encode_tc( + *data_gpu.device_ptr() as *const f64, + state_ptr as *mut c_void, + num_samples, + state_len, + num_qubits as u32, + if self.enable_zz { 1 } else { 0 }, + std::ptr::null_mut(), + ) + }; + + if ret != 0 { + return Err(MahoutError::KernelLaunch(format!( + "Batch IQP TC encoding kernel failed: {} ({})", + ret, + cuda_error_to_string(ret) + ))); + } + } + + { + crate::profile_scope!("GPU::SynchronizeTC"); + device + .synchronize() + .map_err(|e| MahoutError::Cuda(format!("Sync failed: {:?}", e)))?; + } + + Ok(batch_state_vector) + } + + #[cfg(not(target_os = "linux"))] + pub fn encode_batch_tc( + &self, + _device: &Arc, + _batch_data: &[f64], + _num_samples: usize, + _sample_size: usize, + _num_qubits: usize, + ) -> Result { + Err(MahoutError::Cuda( + "CUDA unavailable (non-Linux stub)".to_string(), + )) + } } impl QuantumEncoder for IqpEncoder { diff --git a/qdp/qdp-core/src/lib.rs b/qdp/qdp-core/src/lib.rs index ac5dd5fe94..39c5cf00a6 100644 --- a/qdp/qdp-core/src/lib.rs +++ b/qdp/qdp-core/src/lib.rs @@ -212,6 +212,29 @@ impl QdpEngine { self.encode_batch_for_pipeline(batch_data, num_samples, sample_size, num_qubits, encoding) } + /// Encode a batch of IQP samples via the Ozaki Kronecker Tensor Core path. + #[cfg(target_os = "linux")] + pub fn encode_batch_tc( + &self, + batch_data: &[f64], + num_samples: usize, + sample_size: usize, + num_qubits: usize, + ) -> Result<*mut DLManagedTensor> { + crate::profile_scope!("Mahout::EncodeBatchTC"); + + let encoder = gpu::encodings::IqpEncoder::full(); + let state_vector = encoder.encode_batch_tc( + &self.device, + batch_data, + num_samples, + sample_size, + num_qubits, + )?; + + Ok(state_vector.to_dlpack()) + } + /// Same as [`encode_batch`](Self::encode_batch) with a resolved [`Encoding`] (no string parse). pub(crate) fn encode_batch_for_pipeline( &self, diff --git a/qdp/qdp-kernels/build.rs b/qdp/qdp-kernels/build.rs index def59d6935..16063432b6 100644 --- a/qdp/qdp-kernels/build.rs +++ b/qdp/qdp-kernels/build.rs @@ -26,9 +26,9 @@ use std::env; use std::process::Command; -const DEFAULT_CUBIN_ARCHES: &[&str] = &["75", "80", "86", "89", "90", "100", "120"]; -const DEFAULT_PTX_CANDIDATES: &[&str] = &["120", "100", "90", "89", "86", "80", "75"]; -const LEGACY_FALLBACK_ARCHES: &[&str] = &["75", "80", "86"]; +const DEFAULT_CUBIN_ARCHES: &[&str] = &["80", "86", "89", "90", "100", "120"]; +const DEFAULT_PTX_CANDIDATES: &[&str] = &["120", "100", "90", "89", "86", "80"]; +const LEGACY_FALLBACK_ARCHES: &[&str] = &["80", "86"]; fn add_sm_target(build: &mut cc::Build, arch: &str) { build.flag("-gencode"); @@ -164,6 +164,7 @@ fn main() { println!("cargo:rerun-if-changed=src/angle.cu"); println!("cargo:rerun-if-changed=src/validation.cu"); println!("cargo:rerun-if-changed=src/iqp.cu"); + println!("cargo:rerun-if-changed=src/iqp_tc.cu"); println!("cargo:rerun-if-changed=src/phase.cu"); println!("cargo:rerun-if-env-changed=QDP_NO_CUDA"); println!("cargo:rerun-if-env-changed=QDP_CUDA_ARCH_LIST"); @@ -224,6 +225,8 @@ fn main() { .file("src/angle.cu") .file("src/validation.cu") .file("src/iqp.cu") + .file("src/iqp_tc.cu") + .file("src/ImplicitHadamardOzaki.cu") .file("src/phase.cu") .compile("kernels"); } diff --git a/qdp/qdp-kernels/src/AdaptiveOzaki.h b/qdp/qdp-kernels/src/AdaptiveOzaki.h new file mode 100644 index 0000000000..8e725cc186 --- /dev/null +++ b/qdp/qdp-kernels/src/AdaptiveOzaki.h @@ -0,0 +1,86 @@ +#pragma once + +#include +#include +#include +#include + +namespace ozaki { + +struct PreScanStats { + double dynamic_range; + double sparsity; + int effective_rank; +}; + +enum class ExecutionMode { + Phase22, + Phase23Hetero, + Phase24ExtremeMix, + Phase25Stacking, + Phase26HybridOzaki +}; + +struct OzakiConfig { + double target_epsilon = 1e-10; + size_t vram_limit_mb = 6144; + ExecutionMode mode = ExecutionMode::Phase22; + bool enable_fp64 = true; + bool enable_fp32 = false; // New for Phase 24 + bool enable_fp16 = false; // New for Phase 24 + bool enable_residual = true; + bool enable_partial_residual = false; + bool enable_kernel_fp64 = false; + int split_fp64_bits = 8; + int split_fp32_bits = 0; // New for Phase 24 + int split_fp16_bits = 0; // New for Phase 24 + int split_tc_bits = 32; + int split_residual_bits = 13; + int warp_fp64 = 2; + int warp_tc = 4; + int warp_residual = 2; + double tile_error_guard = 1e-9; +}; + +class AdaptiveOzakiEngine { +public: + explicit AdaptiveOzakiEngine(const OzakiConfig& config); + ~AdaptiveOzakiEngine(); + + void execute(const double* d_A, const double* d_B, double* d_C, int m, int n, int k); + + // Workspace management + void allocateWorkspace(int m, int n, int k); + void freeWorkspace(); + + static void bitShiftToINT8(const double* d_src, int8_t* d_dst_int8, int elem_count, int exponent_shift); + static void computeResidualGradedRing(const double* d_A, const double* d_B, const double* d_C, double* d_R, int m, int n, int k); + static void accumulateSlicedProduct(const int8_t* d_A_int8, const int8_t* d_B_int8, double* d_C, int m, int n, int k, int total_shift); + static void accumulateFusedSlicedProductWMMA(const double* d_A, const double* d_B, double* d_C, int m, int n, int k, int shift_A, int shift_B); + +protected: + PreScanStats analyzeMatrix(const double* d_A, const double* d_B, int m, int n, int k); + int calculateOptimalTile(int m, int n, int k); + +private: + void profileHardwareRatio(); + + OzakiConfig config_; + double ratio_fp64_tc = -1.0; // -1 means not profiled yet + double ratio_fp32_tc = -1.0; + double ratio_fp16_tc = -1.0; + double ratio_tf32_tc = -1.0; + double ratio_int32_tc = -1.0; + + // Workspace pointers + int8_t *dA8_h = nullptr, *dA8_l = nullptr, *dB8_h = nullptr, *dB8_l = nullptr; + uint64_t *dmA_h = nullptr, *dmA_l = nullptr, *dmB_h = nullptr, *dmB_l = nullptr; + int* d_global_work_queue = nullptr; + + // Low-precision buffers for cross-terms + float *dA_hi_f32 = nullptr, *dA_low_f32 = nullptr; + float *dB_hi_f32 = nullptr, *dB_low_f32 = nullptr; + + bool workspace_allocated_ = false; +}; +} diff --git a/qdp/qdp-kernels/src/ImplicitHadamardOzaki.cu b/qdp/qdp-kernels/src/ImplicitHadamardOzaki.cu new file mode 100644 index 0000000000..cd3a1c4ab6 --- /dev/null +++ b/qdp/qdp-kernels/src/ImplicitHadamardOzaki.cu @@ -0,0 +1,330 @@ +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "ImplicitHadamardOzaki.h" +#include +#include +#include +#include +#include + +using namespace nvcuda; + +namespace implicit_ozaki_kernels { + +__device__ __forceinline__ void mma_m16n8k32_s8( + int32_t* d, const uint32_t* a, const uint32_t* b, const int32_t* c) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), + "r"(b[0]), "r"(b[1]), + "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]) + ); +} + +__device__ __forceinline__ void ldmatrix_x4_int8(uint32_t* d, void* smem_ptr) { + uint32_t smem_addr = __cvta_generic_to_shared(smem_ptr); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3]) : "r"(smem_addr)); +} + +__device__ __forceinline__ void ldmatrix_x2_int8(uint32_t* d, void* smem_ptr) { + uint32_t smem_addr = __cvta_generic_to_shared(smem_ptr); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];" + : "=r"(d[0]), "=r"(d[1]) : "r"(smem_addr)); +} + +__device__ __forceinline__ size_t get_A8_offset(int r, int c, int m, int k) { + int tile_m = r / 128; + int tile_k = c / 32; + int num_tiles_k = (k + 31) / 32; + int in_tile_r = r % 128; + int in_tile_c = c % 32; + return (size_t)(tile_m * num_tiles_k + tile_k) * 4096 + in_tile_r * 32 + in_tile_c; +} + +__global__ void precompute_modulo_kernel_p26_implicit(const double* __restrict__ s, int8_t* __restrict__ d, int r, int c, int m, double sh, double sl) { + int local_k = threadIdx.x; int local_m = threadIdx.y; + int tile_k = blockIdx.x; int tile_m = blockIdx.y; + const int pr[7] = {127, 113, 109, 107, 103, 101, 97}; + size_t padded_size = (size_t)((r + 127) / 128) * ((c + 31) / 32) * 4096; + int num_tiles_k = (c + 31) / 32; + int m_idx = (tile_m / 4) * 128 + (tile_m % 4) * 32 + local_m; + int k_idx = tile_k * 32 + local_k; + if (m_idx < r && k_idx < c) { + double v = s[(size_t)m_idx * c + k_idx]; + int32_t iv = (v == 0.0) ? 0 : __double2int_rn(v * sh); + size_t out_off = (size_t)( (m_idx / 128) * num_tiles_k + tile_k ) * 4096 + (m_idx % 128) * 32 + local_k; + for (int p = 0; p < 7; p++) { + int32_t rem = iv % pr[p]; + if (rem < 0) rem += pr[p]; + d[p * padded_size + out_off] = (int8_t)rem; + } + } +} + +template +__device__ void implicit_ozaki_process_one_tile( + int ct, + const int8_t* __restrict__ A8_h, + double* __restrict__ C, + int m, int n, int k, + double inv, + double norm_factor, + int8_t* sA8, + int8_t* sB8, + const int8_t* h_pos, + const int8_t* h_neg, + int warp_id, + int lane_id) { + + constexpr int kTileBytes = 2048; + constexpr int kBufferCount = SingleBuffer ? 1 : 2; + constexpr int kPrimeStride = kTileBytes * kBufferCount; + + const int pr[7] = {127, 113, 109, 107, 103, 101, 97}; + const uint64_t M = 168897325606883ULL; + const uint64_t f[7] = { + 147618922380819ULL, 112099994871825ULL, 134807957135769ULL, + 34726552928518ULL, 96747011755399ULL, 130435558389474ULL, + 19153304965729ULL + }; + + const int total_tiles_m = (m + 63) / 64; + const int tile_m = (ct % total_tiles_m) * 64; + const int tile_n = (ct / total_tiles_m) * 64; + + int32_t prime_acc[7][4][8][4]; + for (int p = 0; p < 7; p++) + for (int i = 0; i < 4; i++) + for (int j = 0; j < 8; j++) + for (int r = 0; r < 4; r++) + prime_acc[p][i][j][r] = 0; + + for (int kk = 0; kk < k; kk += 32) { + const int b_idx = SingleBuffer ? 0 : ((kk / 32) & 1); + const size_t padded_mk = (size_t)((m + 127) / 128) * ((k + 31) / 32) * 4096; + + for (int p = 0; p < 7; ++p) { + const int8_t* Ap_global = A8_h + p * padded_mk; + int8_t* sA_p = &sA8[p * kPrimeStride + b_idx * kTileBytes]; + int8_t* sB_p = &sB8[p * kPrimeStride + b_idx * kTileBytes]; + + for (int i = threadIdx.x; i < 64 * 32; i += blockDim.x) { + int r = i / 32, c = i % 32; + sA_p[r * 32 + c] = (tile_m + r < m && kk + c < k) + ? Ap_global[get_A8_offset(tile_m + r, kk + c, m, k)] : 0; + } + for (int i = threadIdx.x; i < 32 * 64; i += blockDim.x) { + int r = i / 64, c = i % 64; + int8_t val = 0; + if (kk + r < k && tile_n + c < n) { + int parity = __popcll((kk + r) & (tile_n + c)) & 1; + val = (parity == 0) ? h_pos[p] : h_neg[p]; + } + sB_p[r * 64 + c] = val; + } + } + __syncthreads(); + + const int wr = warp_id / 4; + const int wc = warp_id % 4; + + for (int p = 0; p < 7; ++p) { + int8_t* sA_p = &sA8[p * 4096 + b_idx * 2048]; + int8_t* sB_p = &sB8[p * 4096 + b_idx * 2048]; + + for (int mt = 0; mt < 2; mt++) { + uint32_t af[4]; + int row_a = wr * 32 + mt * 16 + (lane_id % 16); + int col_a = (lane_id / 16) * 16; + ldmatrix_x4_int8(af, &sA_p[row_a * 32 + col_a]); + + uint32_t bf_all[4]; + int row_b = (lane_id % 16); + int col_b = wc * 16; + ldmatrix_x4_int8(bf_all, &sB_p[row_b * 64 + col_b]); + + for (int nt = 0; nt < 2; nt++) { + uint32_t bf[2]; + bf[0] = bf_all[nt * 2]; + bf[1] = bf_all[nt * 2 + 1]; + + mma_m16n8k32_s8( + prime_acc[p][wr * 2 + mt][wc * 2 + nt], af, bf, + prime_acc[p][wr * 2 + mt][wc * 2 + nt]); + } + } + } + __syncthreads(); + } + + const int wr = warp_id / 4; + const int wc = warp_id % 4; + for (int mt = 0; mt < 2; mt++) { + for (int nt = 0; nt < 2; nt++) { + uint64_t final_acc[4] = {0, 0, 0, 0}; + for (int p = 0; p < 7; ++p) { + for (int r = 0; r < 4; r++) { + int32_t v = prime_acc[p][wr * 2 + mt][wc * 2 + nt][r]; + uint32_t rem = (v % pr[p] + pr[p]) % pr[p]; + final_acc[r] += (uint64_t)rem * f[p]; + } + } + + int r_base = tile_m + wr * 32 + mt * 16 + (lane_id % 8); + int c_base = tile_n + wc * 16 + nt * 8 + (lane_id / 8) * 2; + + auto store_res = [&](int r, int c, uint64_t cv) { + if (r < m && c < n) { + double fv = (double)(cv % M); + if (cv % M > M / 2) fv -= (double)M; + C[(size_t)r * n + c] = fv * norm_factor * inv; + } + }; + + store_res(r_base, c_base, final_acc[0]); + store_res(r_base + 8, c_base, final_acc[1]); + store_res(r_base, c_base + 1, final_acc[2]); + store_res(r_base + 8, c_base + 1, final_acc[3]); + } + } +} + +__device__ void implicit_ozaki_init_hadamard_signs(int8_t* h_pos, int8_t* h_neg) { + const int pr[7] = {127, 113, 109, 107, 103, 101, 97}; + if (threadIdx.x == 0) { + for (int p = 0; p < 7; p++) { + h_pos[p] = 1 % pr[p]; + int rem_neg = (-1) % pr[p]; + if (rem_neg < 0) rem_neg += pr[p]; + h_neg[p] = rem_neg; + } + } +} + +// NCU / WDDM: one tile per block, static shared memory (28 KiB), single-buffer MMA path. +__global__ void implicit_hadamard_ozaki_grid_kernel_implicit( + const int8_t* __restrict__ A8_h, + double* __restrict__ C, int m, int n, int k, double inv, + double norm_factor) { + + constexpr int kS8Bytes = 7 * 2048; + __shared__ alignas(16) int8_t shared_mem[2 * kS8Bytes]; + int8_t* sA8 = &shared_mem[0]; + int8_t* sB8 = &shared_mem[kS8Bytes]; + + __shared__ int8_t h_pos[7]; + __shared__ int8_t h_neg[7]; + implicit_ozaki_init_hadamard_signs(h_pos, h_neg); + __syncthreads(); + + const int ct = blockIdx.x; + const int total_tiles = ((m + 63) / 64) * ((n + 63) / 64); + if (ct >= total_tiles) return; + + const int warp_id = threadIdx.x / 32; + const int lane_id = threadIdx.x % 32; + implicit_ozaki_process_one_tile( + ct, A8_h, C, m, n, k, inv, norm_factor, + sA8, sB8, h_pos, h_neg, warp_id, lane_id); +} + +template +__global__ void implicit_hadamard_ozaki_persistent_kernel_implicit( + const int8_t* __restrict__ A8_h, + double* __restrict__ C, int m, int n, int k, double inv, int* __restrict__ d_work_queue, + double norm_factor) { + + constexpr int kTileBytes = 2048; + constexpr int kBufferCount = SingleBuffer ? 1 : 2; + constexpr int kPrimeStride = kTileBytes * kBufferCount; + constexpr int kS8Bytes = 7 * kPrimeStride; + extern __shared__ int8_t shared_mem[]; + int8_t* sA8 = &shared_mem[0]; + int8_t* sB8 = &shared_mem[kS8Bytes]; + + const int warp_id = threadIdx.x / 32; + const int lane_id = threadIdx.x % 32; + __shared__ int tile_idx_shared; + + const int total_tiles = ((m + 63) / 64) * ((n + 63) / 64); + + __shared__ int8_t h_pos[7]; + __shared__ int8_t h_neg[7]; + implicit_ozaki_init_hadamard_signs(h_pos, h_neg); + __syncthreads(); + + while (true) { + if (threadIdx.x == 0) tile_idx_shared = atomicAdd(d_work_queue, 1); + __syncthreads(); + const int ct = tile_idx_shared; + if (ct >= total_tiles) break; + + implicit_ozaki_process_one_tile( + ct, A8_h, C, m, n, k, inv, norm_factor, + sA8, sB8, h_pos, h_neg, warp_id, lane_id); + __syncthreads(); + } +} + +} // namespace implicit_ozaki_kernels + +namespace ozaki { + +void ImplicitHadamardOzakiEngine::execute_implicit_hadamard(const double* d_A, double* d_C, int m, int n, int k, double norm_factor, cudaStream_t stream) { + size_t padded_mk = (size_t)((m + 127) / 128) * ((k + 31) / 32) * 4096; + + int8_t *dA8_h = nullptr; + int *d_queue = nullptr; + + cudaMalloc(&dA8_h, 7ULL * padded_mk); + + double scale_A = pow(2.0, 30.0); + double inv_A = pow(2.0, -30.0); + + dim3 pre_block(32, 32); + dim3 pre_grid_A((k + 31) / 32, ((m + 127) / 128) * 4); + implicit_ozaki_kernels::precompute_modulo_kernel_p26_implicit<<>>(d_A, dA8_h, m, k, 0, scale_A, 0.0); + + const bool ncu_profile = (std::getenv("OZAKI_NCU_PROFILE") != nullptr); + if (ncu_profile) { + // Profiling path: grid kernel (one tile per block), single-buffer smem — NCU-friendly on WDDM. + const int total_tiles = ((m + 63) / 64) * ((n + 63) / 64); + implicit_ozaki_kernels::implicit_hadamard_ozaki_grid_kernel_implicit<<>>( + dA8_h, d_C, m, n, k, inv_A, norm_factor); + } else { + int num_sms; + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0); + cudaMalloc(&d_queue, sizeof(int)); + cudaMemsetAsync(d_queue, 0, sizeof(int), stream); + cudaFuncSetAttribute( + implicit_ozaki_kernels::implicit_hadamard_ozaki_persistent_kernel_implicit, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 60 * 1024); + implicit_ozaki_kernels::implicit_hadamard_ozaki_persistent_kernel_implicit<<>>( + dA8_h, d_C, m, n, k, inv_A, d_queue, norm_factor); + } + + cudaStreamSynchronize(stream); + + cudaFree(dA8_h); + if (d_queue) cudaFree(d_queue); +} + +} // namespace ozaki diff --git a/qdp/qdp-kernels/src/ImplicitHadamardOzaki.h b/qdp/qdp-kernels/src/ImplicitHadamardOzaki.h new file mode 100644 index 0000000000..0df10fb6ad --- /dev/null +++ b/qdp/qdp-kernels/src/ImplicitHadamardOzaki.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include +#include +#include + +#include "AdaptiveOzaki.h" + +namespace ozaki { + +class ImplicitHadamardOzakiEngine { +public: + explicit ImplicitHadamardOzakiEngine(const OzakiConfig& config) : config_(config) {} + ~ImplicitHadamardOzakiEngine() = default; + + void execute_implicit_hadamard(const double* d_A, double* d_C, int m, int n, int k, double norm_factor, cudaStream_t stream = 0); + +private: + OzakiConfig config_; +}; + +} // namespace ozaki diff --git a/qdp/qdp-kernels/src/amplitude.cu b/qdp/qdp-kernels/src/amplitude.cu index 57fa4320cf..e434c8e3db 100644 --- a/qdp/qdp-kernels/src/amplitude.cu +++ b/qdp/qdp-kernels/src/amplitude.cu @@ -41,27 +41,18 @@ __global__ void amplitude_encode_kernel( double v1 = 0.0; double v2 = 0.0; - // Vectorized Load Optimization: - // If we are well within bounds, treat input as double2 to issue a single 128-bit load instruction. - // Use __ldg() to pull through the read-only cache; cudaMalloc aligns to 256 bytes so the - // reinterpret_cast load is naturally aligned. + // double2 load via __ldg when aligned and in bounds. if (state_idx_base + 1 < input_len) { - // Reinterpret cast to load two doubles at once const double2 loaded = __ldg(reinterpret_cast(input) + idx); v1 = loaded.x; v2 = loaded.y; } - // Handle edge case: Odd input length else if (state_idx_base < input_len) { v1 = __ldg(input + state_idx_base); - // v2 remains 0.0 } - // Write output: - // Apply pre-calculated reciprocal (multiplication is faster than division) state[state_idx_base] = make_cuDoubleComplex(v1 * inv_norm, 0.0); - // Check boundary for the second element (state_len is usually power of 2, but good to be safe) if (state_idx_base + 1 < state_len) { state[state_idx_base + 1] = make_cuDoubleComplex(v2 * inv_norm, 0.0); } @@ -82,7 +73,6 @@ __global__ void amplitude_encode_kernel_f32( float v2 = 0.0f; if (state_idx_base + 1 < input_len) { - // Mirror the double kernel: cached vectorized load for two floats const float2 loaded = __ldg(reinterpret_cast(input) + idx); v1 = loaded.x; v2 = loaded.y; @@ -225,17 +215,7 @@ int launch_amplitude_encode_f32( return (int)cudaGetLastError(); } -/// Optimized batch amplitude encoding kernel -/// -/// Memory Layout (row-major): -/// - input_batch: [sample0_data | sample1_data | ... | sampleN_data] -/// - state_batch: [sample0_state | sample1_state | ... | sampleN_state] -/// -/// Optimizations: -/// 1. Vectorized double2 loads for 128-bit memory transactions when aligned -/// 2. Grid-stride loop for arbitrary batch sizes -/// 3. Coalesced memory access within warps -/// 4. Scalar fallback for misaligned sample bases and odd tails +/// Batch amplitude encoding kernel (grid-stride, vectorized loads when aligned). __global__ void amplitude_encode_batch_kernel( const double* __restrict__ input_batch, cuDoubleComplex* __restrict__ state_batch, @@ -244,25 +224,20 @@ __global__ void amplitude_encode_batch_kernel( size_t input_len, size_t state_len ) { - // Grid-stride loop pattern for flexibility - const size_t elements_per_sample = state_len / 2; // Each thread handles 2 elements + const size_t elements_per_sample = state_len / 2; const size_t total_work = num_samples * elements_per_sample; const size_t stride = gridDim.x * blockDim.x; size_t global_idx = blockIdx.x * blockDim.x + threadIdx.x; - // Process elements in grid-stride fashion for (size_t idx = global_idx; idx < total_work; idx += stride) { - // Decompose linear index into (sample, element_pair) const size_t sample_idx = idx / elements_per_sample; const size_t elem_pair = idx % elements_per_sample; - // Calculate base addresses (strength-reduced) const size_t input_base = sample_idx * input_len; const size_t state_base = sample_idx * state_len; const size_t elem_offset = elem_pair * 2; - // Load inverse norm (cached by L1) const double inv_norm = inv_norms[sample_idx]; double v1, v2; @@ -281,16 +256,12 @@ __global__ void amplitude_encode_batch_kernel( ? __ldg(sample_input + elem_offset + 1) : 0.0; } else { - // Padding region v1 = v2 = 0.0; } - // Normalize and write as complex numbers - // Compiler will optimize multiplications const cuDoubleComplex c1 = make_cuDoubleComplex(v1 * inv_norm, 0.0); const cuDoubleComplex c2 = make_cuDoubleComplex(v2 * inv_norm, 0.0); - // Write to global memory (coalesced within warp) state_batch[state_base + elem_offset] = c1; if (elem_offset + 1 < state_len) { state_batch[state_base + elem_offset + 1] = c2; @@ -298,17 +269,7 @@ __global__ void amplitude_encode_batch_kernel( } } -/// Optimized batch amplitude encoding kernel (float32) -/// -/// Memory Layout (row-major): -/// - input_batch: [sample0_data | sample1_data | ... | sampleN_data] -/// - state_batch: [sample0_state | sample1_state | ... | sampleN_state] -/// -/// Optimizations: -/// 1. Vectorized float2 loads for 64-bit memory transactions -/// 2. Grid-stride loop for arbitrary batch sizes -/// 3. Coalesced memory access within warps -/// 4. Minimized register pressure +/// Batch amplitude encoding kernel (float32). __global__ void amplitude_encode_batch_kernel_f32( const float* __restrict__ input_batch, cuComplex* __restrict__ state_batch, @@ -317,25 +278,20 @@ __global__ void amplitude_encode_batch_kernel_f32( size_t input_len, size_t state_len ) { - // Grid-stride loop pattern for flexibility const size_t elements_per_sample = state_len / 2; const size_t total_work = num_samples * elements_per_sample; const size_t stride = gridDim.x * blockDim.x; size_t global_idx = blockIdx.x * blockDim.x + threadIdx.x; - // Process elements in grid-stride fashion for (size_t idx = global_idx; idx < total_work; idx += stride) { - // Decompose linear index into (sample, element_pair) const size_t sample_idx = idx / elements_per_sample; const size_t elem_pair = idx % elements_per_sample; - // Calculate base addresses (strength-reduced) const size_t input_base = sample_idx * input_len; const size_t state_base = sample_idx * state_len; const size_t elem_offset = elem_pair * 2; - // Load inverse norm (cached by L1) const float inv_norm = inv_norms[sample_idx]; float v1, v2; @@ -357,11 +313,9 @@ __global__ void amplitude_encode_batch_kernel_f32( v1 = v2 = 0.0f; } - // Normalize and write as complex numbers const cuComplex c1 = make_cuComplex(v1 * inv_norm, 0.0f); const cuComplex c2 = make_cuComplex(v2 * inv_norm, 0.0f); - // Write to global memory (coalesced within warp) state_batch[state_base + elem_offset] = c1; if (elem_offset + 1 < state_len) { state_batch[state_base + elem_offset + 1] = c2; @@ -397,14 +351,9 @@ int launch_amplitude_encode_batch( cuDoubleComplex* state_complex_d = static_cast(state_batch_d); - // Optimal configuration for modern GPUs (SM 7.0+) - // - Block size: DEFAULT_BLOCK_SIZE threads (8 warps, good occupancy) - // - Grid size: Enough blocks to saturate GPU, but not excessive const int blockSize = DEFAULT_BLOCK_SIZE; const size_t total_work = num_samples * (state_len / 2); - // Calculate grid size: aim for high occupancy without too many blocks - // Limit to reasonable number of blocks to avoid scheduler overhead const size_t blocks_needed = (total_work + blockSize - 1) / blockSize; const size_t max_blocks = MAX_GRID_BLOCKS; const size_t gridSize = (blocks_needed < max_blocks) ? blocks_needed : max_blocks; @@ -462,7 +411,6 @@ __global__ void l2_norm_kernel( size_t input_len, double* __restrict__ out_accum ) { - // Vectorized double2 loads for bandwidth and coalescing const size_t vec_idx = blockIdx.x * blockDim.x + threadIdx.x; const size_t stride = gridDim.x * blockDim.x; @@ -497,7 +445,6 @@ __global__ void l2_norm_kernel_f32( size_t input_len, float* __restrict__ out_accum ) { - // Vectorized float2 loads for bandwidth and coalescing const size_t vec_idx = blockIdx.x * blockDim.x + threadIdx.x; const size_t stride = gridDim.x * blockDim.x; diff --git a/qdp/qdp-kernels/src/iqp.cu b/qdp/qdp-kernels/src/iqp.cu index 51fd022ff7..8a23443b2f 100644 --- a/qdp/qdp-kernels/src/iqp.cu +++ b/qdp/qdp-kernels/src/iqp.cu @@ -86,10 +86,7 @@ __device__ cuDoubleComplex compute_amplitude_naive( return make_cuDoubleComplex(real_sum, imag_sum); } -// ============================================================================ -// Naive Implementation: O(2^n) per amplitude, O(4^n) for the full state -// (kept as fallback for small n and verification) -// ============================================================================ +// Naive O(2^n) per amplitude; fallback when FWT overhead dominates. __global__ void iqp_encode_kernel_naive( const double* __restrict__ data, @@ -109,9 +106,7 @@ __global__ void iqp_encode_kernel_naive( } -// ============================================================================ -// FWT O(n * 2^n) Implementation -// ============================================================================ +// FWT O(n * 2^n) path. // Step 1: Compute f[x] = exp(i*theta(x)) for all x. // Uses a grid-stride loop so large state vectors can reuse a fixed launch size. @@ -259,9 +254,7 @@ __global__ void normalize_state_kernel( } } -// ============================================================================ -// Naive O(4^n) Batch Implementation (kept as fallback) -// ============================================================================ +// Naive batch fallback for small n. __global__ void iqp_encode_batch_kernel_naive( const double* __restrict__ data_batch, @@ -275,7 +268,6 @@ __global__ void iqp_encode_batch_kernel_naive( const size_t total_elements = num_samples * state_len; const size_t stride = gridDim.x * blockDim.x; const size_t state_mask = state_len - 1; - // Normalize by 1/2^n (state_len = 2^n) - hoisted outside the loop const double norm = 1.0 / (double)state_len; for (size_t global_idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -292,9 +284,7 @@ __global__ void iqp_encode_batch_kernel_naive( } -// ============================================================================ -// FWT O(n * 2^n) Batch Implementation -// ============================================================================ +// FWT batch path. // Step 1: Compute the normalized phase vector for all samples in batch. __global__ void iqp_phase_batch_kernel( diff --git a/qdp/qdp-kernels/src/iqp_tc.cu b/qdp/qdp-kernels/src/iqp_tc.cu new file mode 100644 index 0000000000..1f905efa97 --- /dev/null +++ b/qdp/qdp-kernels/src/iqp_tc.cu @@ -0,0 +1,289 @@ +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// iqp_tc.cu +#include +#include +#include +#include "kernel_config.h" + +// Phase computation (from iqp.cu) +__device__ double compute_phase_tc( + const double* __restrict__ data, + size_t x, + unsigned int num_qubits, + int enable_zz +) { + double phase = 0.0; + for (unsigned int i = 0; i < num_qubits; ++i) { + phase += data[i] * (double)((x >> i) & 1U); + } + if (enable_zz) { + unsigned int pair_idx = num_qubits; + for (unsigned int i = 0; i < num_qubits; ++i) { + for (unsigned int j = i + 1; j < num_qubits; ++j) { + phase += data[pair_idx] * (double)(((x >> i) & 1U) & ((x >> j) & 1U)); + pair_idx++; + } + } + } + return phase; +} + +// PR3: Shared-memory FWT path (Operator Fusion) +// Fuses Phase computation, Fast Walsh-Hadamard Transform, and Normalization +// entirely within Shared Memory. This completely avoids DRAM roundtrips for N <= 12. +__global__ void iqp_phase_fwt_normalize_tc_kernel( + const double* __restrict__ data_batch, + cuDoubleComplex* __restrict__ state_batch, + size_t num_samples, + size_t state_len, + unsigned int num_qubits, + unsigned int data_len, + int enable_zz, + double norm_factor +) { + extern __shared__ cuDoubleComplex shared_state[]; + + size_t tid = threadIdx.x; + size_t sample_idx = blockIdx.x; + + if (sample_idx >= num_samples) return; + + const double* data = data_batch + sample_idx * data_len; + cuDoubleComplex* state = state_batch + sample_idx * state_len; + + // 1. Phase calculation directly into Shared Memory + for (size_t i = tid; i < state_len; i += blockDim.x) { + double phase = compute_phase_tc(data, i, num_qubits, enable_zz); + double cos_phase, sin_phase; + sincos(phase, &sin_phase, &cos_phase); + shared_state[i] = make_cuDoubleComplex(cos_phase, sin_phase); + } + __syncthreads(); + + // 2. Perform Hadamard FWT in Shared Memory + for (unsigned int stage = 0; stage < num_qubits; ++stage) { + size_t stride = 1ULL << stage; + size_t block_size = stride << 1; + size_t num_pairs = state_len >> 1; + + for (size_t pair_idx = tid; pair_idx < num_pairs; pair_idx += blockDim.x) { + size_t block_idx = pair_idx / stride; + size_t pair_offset = pair_idx % stride; + size_t i = block_idx * block_size + pair_offset; + size_t j = i + stride; + + cuDoubleComplex a = shared_state[i]; + cuDoubleComplex b = shared_state[j]; + + shared_state[i] = cuCadd(a, b); + shared_state[j] = cuCsub(a, b); + } + __syncthreads(); + } + + // 3. Normalize and write back to Global Memory + for (size_t i = tid; i < state_len; i += blockDim.x) { + cuDoubleComplex val = shared_state[i]; + state[i] = make_cuDoubleComplex( + cuCreal(val) * norm_factor, + cuCimag(val) * norm_factor + ); + } +} + +// PR2: Pre-GEMM setup - Unroll Batch and compute initial Phase (split into pure real/imaginary parts) +// This prepares the data layout for the Kronecker product decomposition in upcoming PRs. +__global__ void iqp_phase_split_kernel( + const double* __restrict__ data_batch, + double* __restrict__ state_real, + double* __restrict__ state_imag, + size_t num_samples, + size_t state_len, + unsigned int num_qubits, + unsigned int data_len, + int enable_zz +) { + const size_t total_elements = num_samples * state_len; + const size_t stride = gridDim.x * blockDim.x; + const size_t state_mask = state_len - 1; + + for (size_t global_idx = blockIdx.x * blockDim.x + threadIdx.x; + global_idx < total_elements; + global_idx += stride) { + const size_t sample_idx = global_idx >> num_qubits; + const size_t x = global_idx & state_mask; + const double* data = data_batch + sample_idx * data_len; + + double phase = compute_phase_tc(data, x, num_qubits, enable_zz); + double cos_phase, sin_phase; + sincos(phase, &sin_phase, &cos_phase); + + state_real[global_idx] = cos_phase; + state_imag[global_idx] = sin_phase; + } +} + +#define TRANSPOSE_TILE_DIM 32 +#define TRANSPOSE_BLOCK_ROWS 8 + +// PR2: Shared Memory Bank-Conflict-Free Batch Transpose +// Essential for reordering the data efficiently before/after Tensor Core FWT matrix multiplications. +__global__ void iqp_tc_batch_transpose_kernel(const double* __restrict__ in, double* __restrict__ out, int B, int rows, int cols) { + // TILE_DIM x (TILE_DIM+1) pad to avoid shared memory bank conflicts + __shared__ double tile[TRANSPOSE_TILE_DIM][TRANSPOSE_TILE_DIM + 1]; + + int b = blockIdx.z; + int x = blockIdx.x * TRANSPOSE_TILE_DIM + threadIdx.x; + int y = blockIdx.y * TRANSPOSE_TILE_DIM + threadIdx.y; + + // Load from global memory (coalesced) into shared memory + for (int j = 0; j < TRANSPOSE_TILE_DIM; j += TRANSPOSE_BLOCK_ROWS) { + if (x < cols && (y + j) < rows) { + tile[threadIdx.y + j][threadIdx.x] = in[b * rows * cols + (y + j) * cols + x]; + } + } + + __syncthreads(); + + // Transposed block coordinates + x = blockIdx.y * TRANSPOSE_TILE_DIM + threadIdx.x; + y = blockIdx.x * TRANSPOSE_TILE_DIM + threadIdx.y; + + // Store from shared memory to global memory (coalesced) + for (int j = 0; j < TRANSPOSE_TILE_DIM; j += TRANSPOSE_BLOCK_ROWS) { + if (x < rows && (y + j) < cols) { + out[b * rows * cols + (y + j) * rows + x] = tile[threadIdx.x][threadIdx.y + j]; + } + } +} + +void iqp_tc_launch_transpose(const double* d_in, double* d_out, int B, int rows, int cols, cudaStream_t stream) { + dim3 block(TRANSPOSE_TILE_DIM, TRANSPOSE_BLOCK_ROWS, 1); + dim3 grid((cols + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM, + (rows + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM, B); + iqp_tc_batch_transpose_kernel<<>>(d_in, d_out, B, rows, cols); +} + +// Implicit Hadamard Engine for Tensor Core Blocked FWT +#include "ImplicitHadamardOzaki.h" + +// GEMM 結果重新組合回 cuDoubleComplex +// This restores the memory layout after Tensor Core matrix multiplications. +__global__ void recombine_complex_kernel( + const double* __restrict__ real_part, + const double* __restrict__ imag_part, + cuDoubleComplex* __restrict__ out, + size_t total_elements +) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < total_elements) { + out[idx] = make_cuDoubleComplex(real_part[idx], imag_part[idx]); + } +} + +extern "C" int launch_iqp_encode_tc( + const double* data_batch_d, + void* state_batch_d, + size_t num_samples, + size_t state_len, + unsigned int num_qubits, + int enable_zz, + cudaStream_t stream +) { + unsigned int data_len = num_qubits; + if (enable_zz) { + data_len = num_qubits + (num_qubits * (num_qubits - 1)) / 2; + } + + if (num_qubits <= FWT_SHARED_MEM_THRESHOLD) { + // For N <= 12, use the fused Shared Memory FWT kernel + double norm_factor = 1.0 / (double)state_len; + // Request max dynamic shared memory for this kernel + cudaError_t res = cudaFuncSetAttribute(iqp_phase_fwt_normalize_tc_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536); + if (res != cudaSuccess) return (int)res; + + iqp_phase_fwt_normalize_tc_kernel<<>>( + data_batch_d, static_cast(state_batch_d), num_samples, state_len, num_qubits, data_len, enable_zz, norm_factor + ); + } else { + // Blocked TC-FWT (Kronecker Product Decomposition) + size_t m_samples = num_samples; + size_t total_elements = m_samples * state_len; + + int n1 = num_qubits / 2; + int n2 = num_qubits - n1; + int dim1 = 1 << n1; + int dim2 = 1 << n2; + + double *d_state_real, *d_state_imag; + double *d_out_real, *d_out_imag; + double *d_temp_real, *d_temp_imag; + cudaMalloc(&d_state_real, total_elements * sizeof(double)); + cudaMalloc(&d_state_imag, total_elements * sizeof(double)); + cudaMalloc(&d_out_real, total_elements * sizeof(double)); + cudaMalloc(&d_out_imag, total_elements * sizeof(double)); + cudaMalloc(&d_temp_real, total_elements * sizeof(double)); + cudaMalloc(&d_temp_imag, total_elements * sizeof(double)); + + // Initialize Phase (Split Real/Imag) + const size_t blocks = (total_elements + DEFAULT_BLOCK_SIZE - 1) / DEFAULT_BLOCK_SIZE; + iqp_phase_split_kernel<<>>( + data_batch_d, d_state_real, d_state_imag, num_samples, state_len, num_qubits, data_len, enable_zz + ); + + // Implicit Hadamard Engine for Matrix-Free Hadamard Tensor Core execution + ozaki::OzakiConfig config; + ozaki::ImplicitHadamardOzakiEngine engine(config); + double norm_factor = 1.0 / (double)state_len; + + // 3. TC-FWT Step 1: Z = X * H_{n2} (X shape: B*dim1 x dim2) + engine.execute_implicit_hadamard(d_state_real, d_out_real, num_samples * dim1, dim2, dim2, 1.0, stream); + engine.execute_implicit_hadamard(d_state_imag, d_out_imag, num_samples * dim1, dim2, dim2, 1.0, stream); + + // 4. TC-FWT Step 2: Transpose (B, dim1, dim2) -> (B, dim2, dim1) + iqp_tc_launch_transpose(d_out_real, d_temp_real, num_samples, dim1, dim2, stream); + iqp_tc_launch_transpose(d_out_imag, d_temp_imag, num_samples, dim1, dim2, stream); + + // 5. TC-FWT Step 3: Y_T = Z_T * H_{n1} (Z_T shape: B*dim2 x dim1) + engine.execute_implicit_hadamard(d_temp_real, d_out_real, num_samples * dim2, dim1, dim1, norm_factor, stream); + engine.execute_implicit_hadamard(d_temp_imag, d_out_imag, num_samples * dim2, dim1, dim1, norm_factor, stream); + + // 6. TC-FWT Step 4: Transpose back (B, dim2, dim1) -> (B, dim1, dim2) + iqp_tc_launch_transpose(d_out_real, d_temp_real, num_samples, dim2, dim1, stream); + iqp_tc_launch_transpose(d_out_imag, d_temp_imag, num_samples, dim2, dim1, stream); + + + + // Recombine and Write back + recombine_complex_kernel<<>>( + d_temp_real, d_temp_imag, static_cast(state_batch_d), total_elements + ); + + cudaStreamSynchronize(stream); + + cudaFree(d_state_real); + cudaFree(d_state_imag); + cudaFree(d_out_real); + cudaFree(d_out_imag); + cudaFree(d_temp_real); + cudaFree(d_temp_imag); + } + + cudaError_t err = cudaGetLastError(); + return (int)err; + } diff --git a/qdp/qdp-kernels/src/kernel_config.h b/qdp/qdp-kernels/src/kernel_config.h index 9a5708ac31..7b7f479b03 100644 --- a/qdp/qdp-kernels/src/kernel_config.h +++ b/qdp/qdp-kernels/src/kernel_config.h @@ -63,8 +63,8 @@ // Threshold for shared memory FWT optimization // For n <= this threshold, use shared memory FWT (single kernel launch) // For n > threshold, use global memory FWT (multiple kernel launches) -// 10 qubits = 2^10 * 16 bytes (cuDoubleComplex) = 16KB shared memory -#define FWT_SHARED_MEM_THRESHOLD 10 +// 12 qubits = 2^12 * 16 bytes (cuDoubleComplex) = 64KB shared memory +#define FWT_SHARED_MEM_THRESHOLD 12 // Minimum qubits to use FWT optimization (below this, naive is competitive) #define FWT_MIN_QUBITS 4 diff --git a/qdp/qdp-kernels/src/lib.rs b/qdp/qdp-kernels/src/lib.rs index 9a1f832d85..8a96e23e11 100644 --- a/qdp/qdp-kernels/src/lib.rs +++ b/qdp/qdp-kernels/src/lib.rs @@ -376,6 +376,21 @@ unsafe extern "C" { stream: *mut c_void, ) -> i32; + /// Launch IQP encoding kernel using Matrix-Free Implicit Hadamard Tensor Core Engine + /// Returns CUDA error code (0 = success) + /// + /// # Safety + /// Requires valid GPU pointers, must sync before freeing + pub fn launch_iqp_encode_tc( + data_batch_d: *const f64, + state_batch_d: *mut c_void, + num_samples: usize, + state_len: usize, + num_qubits: u32, + enable_zz: i32, + stream: *mut c_void, + ) -> i32; + /// Launch phase encoding kernel /// Returns CUDA error code (0 = success) /// @@ -700,6 +715,20 @@ pub extern "C" fn launch_iqp_encode_batch( 999 } +#[cfg(any(not(target_os = "linux"), qdp_no_cuda))] +#[unsafe(no_mangle)] +pub extern "C" fn launch_iqp_encode_tc( + _data_batch_d: *const f64, + _state_batch_d: *mut c_void, + _num_samples: usize, + _state_len: usize, + _num_qubits: u32, + _enable_zz: i32, + _stream: *mut c_void, +) -> i32 { + 999 +} + #[cfg(any(not(target_os = "linux"), qdp_no_cuda))] #[unsafe(no_mangle)] pub extern "C" fn launch_phase_encode( diff --git a/qdp/qdp-python/qumat_qdp/backend.py b/qdp/qdp-python/qumat_qdp/backend.py index 42f1051720..9dbfb64b55 100644 --- a/qdp/qdp-python/qumat_qdp/backend.py +++ b/qdp/qdp-python/qumat_qdp/backend.py @@ -118,3 +118,16 @@ def encode( :raises ValueError: If the backend does not support ``encoding_method``. """ return self._engine_adapter.encode(data, num_qubits, encoding_method) + + def encode_batch_tc( + self, + data: Any, + num_qubits: int, + ) -> Any: + """Encode a batch of IQP samples via the Ozaki Tensor Core path.""" + encode_tc = getattr(self._engine_adapter, "encode_batch_tc", None) + if encode_tc is None: + raise RuntimeError( + "encode_batch_tc is unavailable. Rebuild with CUDA on Linux/WSL." + ) + return encode_tc(data, num_qubits) diff --git a/qdp/qdp-python/src/engine.rs b/qdp/qdp-python/src/engine.rs index 8297bf5e78..4047abd227 100644 --- a/qdp/qdp-python/src/engine.rs +++ b/qdp/qdp-python/src/engine.rs @@ -127,6 +127,32 @@ impl QdpEngine { self.encode_from_list(data, num_qubits, encoding_method) } + /// Encode a batch of IQP samples using the Ozaki Kronecker Tensor Core path. + #[cfg(target_os = "linux")] + fn encode_batch_tc( + &self, + data: &Bound<'_, PyAny>, + num_qubits: usize, + ) -> PyResult { + let array_2d = data.extract::>().map_err(|_| { + PyRuntimeError::new_err("Failed to extract 2D NumPy array. Ensure dtype is float64.") + })?; + let shape = array_2d.shape(); + let num_samples = shape[0]; + let sample_size = shape[1]; + let data_slice = array_2d + .as_slice() + .map_err(|_| PyRuntimeError::new_err("NumPy array must be contiguous (C-order)"))?; + let ptr = self + .engine + .encode_batch_tc(data_slice, num_samples, sample_size, num_qubits) + .map_err(|e| PyRuntimeError::new_err(format!("Encoding failed: {}", e)))?; + Ok(QuantumTensor { + ptr, + consumed: false, + }) + } + /// Encode from NumPy array (1D or 2D) fn encode_from_numpy( &self, diff --git a/testing/qdp/test_iqp_tc_path.py b/testing/qdp/test_iqp_tc_path.py new file mode 100644 index 0000000000..d35b9d94ce --- /dev/null +++ b/testing/qdp/test_iqp_tc_path.py @@ -0,0 +1,85 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Smoke and normalization tests for FWT vs Ozaki TC IQP paths (GPU vs GPU).""" + +import pytest +import torch +from qumat_qdp import QdpEngine + + +def _iqp_param_count(num_qubits: int) -> int: + return num_qubits + num_qubits * (num_qubits - 1) // 2 + + +@pytest.fixture(scope="module") +def engine(): + try: + eng = QdpEngine(device_id=0, precision="float64") + except Exception as exc: + pytest.skip(f"Could not initialize QdpEngine: {exc}") + if not hasattr(eng, "encode_batch_tc"): + pytest.skip("encode_batch_tc not available in this build") + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + return eng + + +def _assert_normalized(state: torch.Tensor, num_qubits: int, label: str) -> None: + probs = state.abs() ** 2 + if state.ndim == 2: + row_sums = probs.sum(dim=1) + assert torch.allclose(row_sums, torch.ones_like(row_sums), atol=1e-6), ( + f"{label}: batch normalization failed at N={num_qubits}" + ) + else: + assert torch.allclose(probs.sum(), torch.tensor(1.0), atol=1e-6), ( + f"{label}: normalization failed at N={num_qubits}" + ) + + +@pytest.mark.parametrize("num_qubits", [8, 12]) +@pytest.mark.parametrize("batch_size", [4, 32]) +def test_fwt_and_tc_paths_normalized(engine, num_qubits, batch_size): + """For N<=12 both GPU paths return normalized states.""" + data_len = _iqp_param_count(num_qubits) + data = torch.randn(batch_size, data_len, dtype=torch.float64).numpy() + state_len = 1 << num_qubits + + fwt_state = torch.from_dlpack(engine.encode(data, num_qubits, "iqp")) + assert fwt_state.shape == (batch_size, state_len) + _assert_normalized(fwt_state, num_qubits, "FWT") + + tc_state = torch.from_dlpack(engine.encode_batch_tc(data, num_qubits)) + assert tc_state.shape == (batch_size, state_len) + _assert_normalized(tc_state, num_qubits, "TC") + + +@pytest.mark.parametrize("num_qubits", [14]) +@pytest.mark.parametrize("batch_size", [4, 8]) +def test_large_n_tc_path_smoke(engine, num_qubits, batch_size): + """Large-N TC Kronecker path runs; FWT remains normalized baseline.""" + data_len = _iqp_param_count(num_qubits) + data = torch.randn(batch_size, data_len, dtype=torch.float64).numpy() + state_len = 1 << num_qubits + + fwt_state = torch.from_dlpack(engine.encode(data, num_qubits, "iqp")) + assert fwt_state.shape == (batch_size, state_len) + _assert_normalized(fwt_state, num_qubits, "FWT") + + tc_state = torch.from_dlpack(engine.encode_batch_tc(data, num_qubits)) + assert tc_state.shape == (batch_size, state_len) + assert torch.isfinite(tc_state).all()