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..452b0e3fd1 100644 --- a/qdp/qdp-kernels/src/iqp.cu +++ b/qdp/qdp-kernels/src/iqp.cu @@ -60,58 +60,7 @@ __device__ double compute_phase( return phase; } -// Compute the unnormalized amplitude for basis state |z> via a naive O(2^n) sum over x -// (O(4^n) total work if amplitudes for all 2^n basis states z are computed): -// sum_x exp(i*theta(x)) * (-1)^popcount(x AND z) -__device__ cuDoubleComplex compute_amplitude_naive( - const double* __restrict__ data, - size_t z, - size_t state_len, - unsigned int num_qubits, - int enable_zz -) { - double real_sum = 0.0; - double imag_sum = 0.0; - - for (size_t x = 0; x < state_len; ++x) { - double phase = compute_phase(data, x, num_qubits, enable_zz); - int parity = __popcll(x & z) & 1; - double sign = (parity == 0) ? 1.0 : -1.0; - double cos_phase, sin_phase; - sincos(phase, &sin_phase, &cos_phase); - real_sum += sign * cos_phase; - imag_sum += sign * sin_phase; - } - - 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) -// ============================================================================ - -__global__ void iqp_encode_kernel_naive( - const double* __restrict__ data, - cuDoubleComplex* __restrict__ state, - size_t state_len, - unsigned int num_qubits, - int enable_zz -) { - size_t z = blockIdx.x * blockDim.x + threadIdx.x; - if (z >= state_len) return; - - cuDoubleComplex amp = compute_amplitude_naive(data, z, state_len, num_qubits, enable_zz); - - // Normalize by 1/2^n (state_len = 2^n) - double norm = 1.0 / (double)state_len; - state[z] = make_cuDoubleComplex(cuCreal(amp) * norm, cuCimag(amp) * norm); -} - - -// ============================================================================ -// 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. @@ -240,6 +189,65 @@ __global__ void iqp_phase_fwt_shared_normalize_kernel( } } +// Fused phase + shared-memory FWT + normalization for one sample in a batch. +// One CUDA block per sample; avoids intermediate global-memory traffic for N <= 12. +__global__ void iqp_phase_fwt_shared_normalize_batch_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; + + for (size_t i = tid; i < state_len; i += blockDim.x) { + double phase = compute_phase(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(); + + 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(); + } + + 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 + ); + } +} + // Step 3: Normalize the state by 1/state_len (= 1/2^n) __global__ void normalize_state_kernel( cuDoubleComplex* __restrict__ state, @@ -259,42 +267,7 @@ __global__ void normalize_state_kernel( } } -// ============================================================================ -// Naive O(4^n) Batch Implementation (kept as fallback) -// ============================================================================ - -__global__ void iqp_encode_batch_kernel_naive( - 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 -) { - 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; - global_idx < total_elements; - global_idx += stride) { - const size_t sample_idx = global_idx >> num_qubits; - const size_t z = global_idx & state_mask; - const double* data = data_batch + sample_idx * data_len; - - cuDoubleComplex amp = compute_amplitude_naive(data, z, state_len, num_qubits, enable_zz); - - state_batch[global_idx] = make_cuDoubleComplex(cuCreal(amp) * norm, cuCimag(amp) * norm); - } -} - - -// ============================================================================ -// 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( @@ -431,20 +404,7 @@ int launch_iqp_encode( const int blockSize = DEFAULT_BLOCK_SIZE; const double norm_factor = 1.0 / (double)state_len; - // Use naive kernel for small n (FWT overhead not worth it) - if (num_qubits < FWT_MIN_QUBITS) { - const int gridSize = (state_len + blockSize - 1) / blockSize; - iqp_encode_kernel_naive<<>>( - data_d, - state_complex_d, - state_len, - num_qubits, - enable_zz - ); - return (int)cudaGetLastError(); - } - - // FWT-based implementation for larger n + // FWT-based implementation const size_t blocks_needed = (state_len + blockSize - 1) / blockSize; const int gridSize = (int)((blocks_needed < MAX_GRID_BLOCKS) ? blocks_needed : MAX_GRID_BLOCKS); @@ -452,6 +412,11 @@ int launch_iqp_encode( // Shared-memory fast path: phase generation, full FWT, and normalization // happen in a single launch and touch global memory only once. size_t shared_mem_size = state_len * sizeof(cuDoubleComplex); + cudaFuncSetAttribute( + iqp_phase_fwt_shared_normalize_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 65536 + ); iqp_phase_fwt_shared_normalize_kernel<<<1, blockSize, shared_mem_size, stream>>>( data_d, state_complex_d, @@ -530,49 +495,49 @@ int launch_iqp_encode_batch( const size_t gridSize = (blocks_needed < MAX_GRID_BLOCKS) ? blocks_needed : MAX_GRID_BLOCKS; const double norm_factor = 1.0 / (double)state_len; - // Use naive kernel for small n (FWT overhead not worth it) - if (num_qubits < FWT_MIN_QUBITS) { - iqp_encode_batch_kernel_naive<<>>( + if (num_qubits <= FWT_SHARED_MEM_THRESHOLD) { + size_t shared_mem_size = state_len * sizeof(cuDoubleComplex); + cudaFuncSetAttribute( + iqp_phase_fwt_shared_normalize_batch_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 65536 + ); + iqp_phase_fwt_shared_normalize_batch_kernel<<>>( data_batch_d, state_complex_d, num_samples, state_len, num_qubits, data_len, - enable_zz + enable_zz, + norm_factor ); - return (int)cudaGetLastError(); - } - - // FWT-based implementation for larger n - - // Step 1: Compute phase array f[x] = exp(i*theta(x)) for all samples - iqp_phase_batch_kernel<<>>( - data_batch_d, - state_complex_d, - num_samples, - state_len, - num_qubits, - data_len, - enable_zz, - norm_factor - ); - - // Step 2: Apply FWT to all samples (global memory version for batch) - // For batch processing, we always use global memory FWT - // (shared memory would require processing samples one at a time) - const size_t total_pairs = num_samples * (state_len >> 1); - const size_t fwt_blocks_needed = (total_pairs + blockSize - 1) / blockSize; - const size_t fwt_grid_size = (fwt_blocks_needed < MAX_GRID_BLOCKS) ? fwt_blocks_needed : MAX_GRID_BLOCKS; - - for (unsigned int stage = 0; stage < num_qubits; ++stage) { - fwt_butterfly_batch_kernel<<>>( + } else { + // Global-memory FWT for larger qubit counts. + iqp_phase_batch_kernel<<>>( + data_batch_d, state_complex_d, num_samples, state_len, num_qubits, - stage + data_len, + enable_zz, + norm_factor ); + + const size_t total_pairs = num_samples * (state_len >> 1); + const size_t fwt_blocks_needed = (total_pairs + blockSize - 1) / blockSize; + const size_t fwt_grid_size = (fwt_blocks_needed < MAX_GRID_BLOCKS) ? fwt_blocks_needed : MAX_GRID_BLOCKS; + + for (unsigned int stage = 0; stage < num_qubits; ++stage) { + fwt_butterfly_batch_kernel<<>>( + state_complex_d, + num_samples, + state_len, + num_qubits, + stage + ); + } } return (int)cudaGetLastError(); diff --git a/qdp/qdp-kernels/src/kernel_config.h b/qdp/qdp-kernels/src/kernel_config.h index 069f69a3b4..28329ca545 100644 --- a/qdp/qdp-kernels/src/kernel_config.h +++ b/qdp/qdp-kernels/src/kernel_config.h @@ -56,8 +56,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-python/benchmark/benchmark_pr3.py b/qdp/qdp-python/benchmark/benchmark_pr3.py new file mode 100644 index 0000000000..1ab7a583f4 --- /dev/null +++ b/qdp/qdp-python/benchmark/benchmark_pr3.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +# +# 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 shared-memory FWT fusion benchmark (PR3). + +For num_qubits <= FWT_SHARED_MEM_THRESHOLD the batch path fuses phase +computation, shared-memory FWT, and normalization in one kernel per sample. + +Run from repo root:: + + uv run --project qdp/qdp-python python qdp/qdp-python/benchmark/benchmark_pr3.py +""" + +from __future__ import annotations + +import argparse +import time + +import numpy as np +import torch +from qumat_qdp import QdpEngine + +FWT_SHARED_MEM_THRESHOLD = 12 + + +def benchmark_iqp_batch( + num_qubits: int, + num_samples: int, + iters: int = 50, +) -> float: + """Return average batch IQP encode latency in microseconds.""" + data_len = num_qubits + num_qubits * (num_qubits - 1) // 2 + batch_data = np.random.randn(num_samples, data_len).astype(np.float64) + engine = QdpEngine(0) + + for _ in range(5): + _ = engine.encode(batch_data, num_qubits, "iqp") + torch.cuda.synchronize() + + start = time.perf_counter() + for _ in range(iters): + _ = engine.encode(batch_data, num_qubits, "iqp") + torch.cuda.synchronize() + elapsed_us = (time.perf_counter() - start) / iters * 1e6 + return elapsed_us + + +def main() -> None: + parser = argparse.ArgumentParser(description="IQP shared-memory FWT benchmark") + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument( + "--qubits", + type=int, + nargs="+", + default=[8, 10, 12, 13, 14], + help="Qubit counts to benchmark", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA not available. Cannot benchmark.") + + device_name = torch.cuda.get_device_name(0) + print("=" * 70) + print("IQP shared-memory FWT fusion benchmark (PR3)") + print(f"GPU: {device_name}") + print( + f"Config: batch_size={args.batch_size}, iterations={args.iterations}, " + f"threshold={FWT_SHARED_MEM_THRESHOLD}" + ) + print("=" * 70) + print(f"{'Qubits':<8} {'Path':<22} {'Time (us)':>12}") + print("-" * 70) + + for n in args.qubits: + path = "shared-mem fused" if n <= FWT_SHARED_MEM_THRESHOLD else "global-mem FWT" + latency_us = benchmark_iqp_batch(n, args.batch_size, args.iterations) + print(f"{n:<8} {path:<22} {latency_us:>12.1f}") + + +if __name__ == "__main__": + main() diff --git a/testing/qdp/test_implicit_fwt.py b/testing/qdp/test_implicit_fwt.py new file mode 100644 index 0000000000..f53735179f --- /dev/null +++ b/testing/qdp/test_implicit_fwt.py @@ -0,0 +1,68 @@ +# +# 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. + +import pytest +import torch +from qumat_qdp import QdpEngine +from qumat_qdp.torch_ref import iqp_encode as iqp_encode_baseline + + +@pytest.fixture(scope="module") +def engine(): + try: + return QdpEngine(precision="float64") + except Exception as e: + pytest.skip(f"Could not initialize QdpEngine: {e}") + + +@pytest.mark.parametrize("n_qubits", [2, 4, 8]) +@pytest.mark.parametrize("batch_size", [1, 16]) +@pytest.mark.parametrize("enable_zz", [True, False]) +def test_implicit_fwt_correctness(engine, n_qubits, batch_size, enable_zz): + """ + Test that the QDP engine's implicit FWT logic perfectly matches + the theoretical exact PyTorch implementation. + """ + # Expected number of parameters for IQP + if enable_zz: + n_params = n_qubits + n_qubits * (n_qubits - 1) // 2 + method = "iqp" + else: + n_params = n_qubits + method = "iqp-z" + + # Generate random parameters + data = torch.randn(batch_size, n_params, dtype=torch.float64, device="cuda") + + # 1. Baseline logic (pure PyTorch) + expected_state = iqp_encode_baseline( + data, n_qubits, enable_zz=enable_zz, device="cuda" + ) + + # 2. QDP implicit FWT logic + # The QDP engine internally dispatches to standard SIMT implicit FWT kernel + actual_state_dlpack = engine.encode(data, n_qubits, encoding_method=method) + actual_state = torch.from_dlpack(actual_state_dlpack) + + # 3. Validation + # We use a strict tolerance since both should be FP64 deterministic computations + torch.testing.assert_close( + actual_state, + expected_state, + rtol=1e-12, + atol=1e-12, + msg=f"Mismatch found for N={n_qubits}, enable_zz={enable_zz}", + ) diff --git a/testing/qdp/test_shared_memory_fwt.py b/testing/qdp/test_shared_memory_fwt.py new file mode 100644 index 0000000000..63d6633a49 --- /dev/null +++ b/testing/qdp/test_shared_memory_fwt.py @@ -0,0 +1,58 @@ +# +# 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. + +import pytest +import torch +from qumat_qdp import QdpEngine +from qumat_qdp.torch_ref import iqp_encode as iqp_encode_baseline + + +@pytest.fixture(scope="module") +def engine(): + try: + return QdpEngine(precision="float64") + except Exception as e: + pytest.skip(f"Could not initialize QdpEngine: {e}") + + +@pytest.mark.parametrize("n_qubits", [8, 10, 12]) +@pytest.mark.parametrize("batch_size", [4, 16, 32]) +@pytest.mark.parametrize("enable_zz", [True, False]) +def test_shared_memory_fwt_batch_correctness(engine, n_qubits, batch_size, enable_zz): + """Validate fused shared-memory batch IQP encoding for N <= 12.""" + if enable_zz: + n_params = n_qubits + n_qubits * (n_qubits - 1) // 2 + method = "iqp" + else: + n_params = n_qubits + method = "iqp-z" + + data = torch.randn(batch_size, n_params, dtype=torch.float64, device="cuda") + + expected_state = iqp_encode_baseline( + data, n_qubits, enable_zz=enable_zz, device="cuda" + ) + + actual_state_dlpack = engine.encode(data, n_qubits, encoding_method=method) + actual_state = torch.from_dlpack(actual_state_dlpack) + + torch.testing.assert_close( + actual_state, + expected_state, + rtol=1e-12, + atol=1e-12, + msg=f"Shared-memory batch mismatch for N={n_qubits}, batch={batch_size}", + )