ryankert01 commented on code in PR #1387: URL: https://github.com/apache/mahout/pull/1387#discussion_r3450819783
########## testing/qdp/test_batch_throughput.py: ########## @@ -0,0 +1,67 @@ +# +# 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 + + [email protected](scope="module") +def engine(): + try: + return QdpEngine(precision="float64") + except Exception as e: + pytest.skip(f"Could not initialize QdpEngine: {e}") + + [email protected]("n_qubits", [2, 4, 6]) [email protected]("batch_size", [1, 16, 64]) [email protected]("enable_zz", [True, False]) +def test_batch_throughput_opt_correctness(engine, n_qubits, batch_size, enable_zz): + """ + Test that the batched IQP logic introduced in PR2 (which splits phase + and transposes for Tensor Cores) yields identical outputs to the theoretical + PyTorch baseline for various batch sizes. + """ + 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 (batched) + data = torch.randn(batch_size, n_params, dtype=torch.float64, device="cuda") + + # 1. Baseline logic (pure PyTorch, O(4^n) equivalent or reference) + expected_state = iqp_encode_baseline( + data, n_qubits, enable_zz=enable_zz, device="cuda" + ) + + # 2. QDP Engine (Testing the Batch throughput C++ / Rust API path) + actual_state_dlpack = engine.encode(data, n_qubits, encoding_method=method) Review Comment: This test currently exercises `engine.encode(...)`, which routes through the existing standard IQP batch path, not the new `encode_batch_tc` path. So it can pass even if the TC scaffold is incorrect. Can we add a direct CUDA/Rust test for `IqpEncoder::encode_batch_tc`, or expose an explicit opt-in Python API and test that path? ########## qdp/qdp-kernels/src/iqp_tc.cu: ########## @@ -0,0 +1,165 @@ +// +// 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 <cuda_runtime.h> +#include <cuComplex.h> +#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; +} + +// 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 + +// 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<<<grid, block, 0, stream>>>(d_in, d_out, B, rows, cols); +} + +// Recombine Real and Imaginary parts back into 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 +) { + // Scaffold for batch layout manipulation + size_t total_elements = num_samples * state_len; + + double *d_state_real, *d_state_imag; + cudaMalloc(&d_state_real, total_elements * sizeof(double)); + cudaMalloc(&d_state_imag, total_elements * sizeof(double)); + + unsigned int data_len = num_qubits; + const size_t blocks = (total_elements + DEFAULT_BLOCK_SIZE - 1) / DEFAULT_BLOCK_SIZE; + + iqp_phase_split_kernel<<<blocks, DEFAULT_BLOCK_SIZE, 0, stream>>>( + data_batch_d, d_state_real, d_state_imag, num_samples, state_len, num_qubits, data_len, enable_zz + ); + + // In future PRs, Kronecker Transpose and FWT will happen here. + + recombine_complex_kernel<<<blocks, DEFAULT_BLOCK_SIZE, 0, stream>>>( Review Comment: This TC launcher currently writes the pre-FWT phase vector rather than an IQP-encoded state. `iqp_phase_split_kernel` computes `exp(i theta(x))`, and this recombine writes it directly to the output; the FWT + normalization is still marked as future work. For all-zero params, this returns all ones instead of `|0...0>`. Can we either keep this as private scaffolding for now, or complete the transform before exposing it through `encode_batch_tc`? ########## qdp/qdp-kernels/src/iqp_tc.cu: ########## @@ -0,0 +1,165 @@ +// +// 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 <cuda_runtime.h> +#include <cuComplex.h> +#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; +} + +// 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 + +// 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<<<grid, block, 0, stream>>>(d_in, d_out, B, rows, cols); +} + +// Recombine Real and Imaginary parts back into 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 +) { + // Scaffold for batch layout manipulation + size_t total_elements = num_samples * state_len; + + double *d_state_real, *d_state_imag; + cudaMalloc(&d_state_real, total_elements * sizeof(double)); + cudaMalloc(&d_state_imag, total_elements * sizeof(double)); + + unsigned int data_len = num_qubits; Review Comment: For full ZZ batches, this hardcodes `data_len = num_qubits`, but Rust validates `sample_size` as `num_qubits + num_qubits * (num_qubits - 1) / 2`. That means sample 1 starts inside sample 0's ZZ parameters when `enable_zz=true`. Can we pass the actual `sample_size`/data length into `launch_iqp_encode_tc`, matching the existing batch launcher? ########## qdp/qdp-kernels/src/iqp_tc.cu: ########## @@ -0,0 +1,165 @@ +// +// 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 <cuda_runtime.h> +#include <cuComplex.h> +#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; +} + +// 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 + +// 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<<<grid, block, 0, stream>>>(d_in, d_out, B, rows, cols); +} + +// Recombine Real and Imaginary parts back into 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 +) { + // Scaffold for batch layout manipulation + size_t total_elements = num_samples * state_len; + + double *d_state_real, *d_state_imag; + cudaMalloc(&d_state_real, total_elements * sizeof(double)); Review Comment: This launcher ignores both `cudaMalloc` return values and then returns `cudaSuccess` unconditionally. On OOM or launch failure, Rust will report success while the output may be invalid; if the second allocation fails, the first allocation can also leak. Can we return CUDA errors consistently with the existing IQP launchers and free any partial allocations on failure? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
