You are viewing a plain text version of this content. The canonical link for it is here.
Posted to dev@mahout.apache.org by "400Ping (via GitHub)" <gi...@apache.org> on 2025/12/22 13:34:11 UTC
[PR] [QDP] Double-buffered async I/O for read_parquet_batch [mahout]
400Ping opened a new pull request, #751:
URL: https://github.com/apache/mahout/pull/751
### Purpose of PR
<!-- Describe what this PR does. -->
read_parquet_batch blocks -> introduce double-buffered async I/O to overlap CPU read with GPU compute and improve throughput.
### Related Issues or PRs
<!-- Add links to related issues or PRs. -->
<!-- - Closes #123 -->
<!-- - Related to #123 -->
Closes #703
### Changes Made
<!-- Please mark one with an "x" -->
- [ ] Bug fix
- [x] New feature
- [ ] Refactoring
- [ ] Documentation
- [ ] Test
- [ ] CI/CD pipeline
- [ ] Other
### Breaking Changes
<!-- Does this PR introduce a breaking change? -->
- [ ] Yes
- [x] No
### Checklist
<!-- Please mark each item with an "x" when complete -->
<!-- If not all items are complete, please open this as a **Draft PR**.
Once all requirements are met, mark as ready for review. -->
- [ ] Added or updated unit tests for all changes
- [ ] Added or updated documentation for all changes
- [ ] Successfully built and ran all unit tests or manual tests locally
- [ ] PR title follows "MAHOUT-XXX: Brief Description" format (if related to an issue)
- [ ] Code follows ASF guidelines
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3702064623
Done, @rich7420 PTAL
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3703215524
cc @guan404ming @ryankert01
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2656461871
##########
qdp/qdp-core/src/gpu/pipeline.rs:
##########
@@ -171,32 +252,40 @@ where
{
crate::profile_scope!("GPU::AsyncPipeline");
- // 1. Create dual streams for pipeline overlap
- let stream1 = device
- .fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("Failed to create stream 1: {:?}", e)))?;
- let stream2 = device
- .fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("Failed to create stream 2: {:?}", e)))?;
- let streams = [&stream1, &stream2];
+ // Pinned host staging pool sized to the current chunking strategy (double-buffer by default).
+ const CHUNK_SIZE_ELEMENTS: usize = 8 * 1024 * 1024 / std::mem::size_of::<f64>(); // 8MB
+ const PINNED_POOL_SIZE: usize = 2; // double buffering
+ // 1. Create dual streams with per-slot events to coordinate copy -> compute
+ let ctx = PipelineContext::new(device, PINNED_POOL_SIZE)?;
+ let pinned_pool = PinnedBufferPool::new(PINNED_POOL_SIZE, CHUNK_SIZE_ELEMENTS)
+ .map_err(|e| MahoutError::Cuda(format!("Failed to create pinned buffer pool: {}", e)))?;
// 2. Chunk size: 8MB per chunk (balance between overhead and overlap opportunity)
- // TODO: we should tune this dynamically based on the detected GPU model or PCIe bandwidth in the future.
- // Too small = launch overhead dominates, too large = less overlap
- const CHUNK_SIZE_ELEMENTS: usize = 8 * 1024 * 1024 / std::mem::size_of::<f64>(); // 8MB
+ // TODO: tune dynamically based on GPU/PCIe bandwidth.
// 3. Keep temporary buffers alive until all streams complete
// This prevents Rust from dropping them while GPU is still using them
let mut keep_alive_buffers: Vec<CudaSlice<f64>> = Vec::new();
+ // Keep pinned buffers alive until the copy stream has completed their H2D copy
+ let mut in_flight_pinned: Vec<PinnedBufferHandle> = Vec::new();
let mut global_offset = 0;
- // 4. Pipeline loop: alternate between streams for maximum overlap
+ // 4. Pipeline loop: copy on copy stream, compute on compute stream with event handoff
for (chunk_idx, chunk) in host_data.chunks(CHUNK_SIZE_ELEMENTS).enumerate() {
- let current_stream = streams[chunk_idx % 2];
+ let chunk_offset = global_offset;
+ let event_slot = chunk_idx % PINNED_POOL_SIZE;
crate::profile_scope!("GPU::ChunkProcess");
+ if chunk.len() > CHUNK_SIZE_ELEMENTS {
+ return Err(MahoutError::InvalidInput(format!(
+ "Chunk size {} exceeds pinned buffer capacity {}",
+ chunk.len(),
+ CHUNK_SIZE_ELEMENTS
+ )));
+ }
+
Review Comment:
This check is redundant because the iterator `chunks()` will never produce a chunk larger than `CHUNK_SIZE_ELEMENTS`. The check can be removed to simplify the code.
```suggestion
```
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -0,0 +1,141 @@
+//
+// 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.
+
+//! Reusable pool of pinned host buffers for staging Disk → Host → GPU transfers.
+//! Intended for producer/consumer pipelines that need a small, fixed set of
+//! page-locked buffers to avoid repeated cudaHostAlloc / cudaFreeHost.
+
+use std::sync::{Arc, Condvar, Mutex};
+
+use crate::error::{MahoutError, Result};
+use crate::gpu::memory::PinnedHostBuffer;
+
+/// Handle that automatically returns a buffer to the pool on drop.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferHandle {
+ buffer: Option<PinnedHostBuffer>,
+ pool: Arc<PinnedBufferPool>,
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::Deref for PinnedBufferHandle {
+ type Target = PinnedHostBuffer;
+
+ fn deref(&self) -> &Self::Target {
+ self.buffer
+ .as_ref()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::DerefMut for PinnedBufferHandle {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ self.buffer
+ .as_mut()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl Drop for PinnedBufferHandle {
+ fn drop(&mut self) {
+ if let Some(buf) = self.buffer.take() {
+ let mut free = self.pool.free.lock().unwrap();
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
+ }
+}
+
+/// Pool of pinned host buffers sized for a fixed batch shape.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferPool {
+ free: Mutex<Vec<PinnedHostBuffer>>,
+ available_cv: Condvar,
+ capacity: usize,
+ elements_per_buffer: usize,
+}
+
+#[cfg(target_os = "linux")]
+impl PinnedBufferPool {
+ /// Create a pool with `pool_size` pinned buffers, each sized for `elements_per_buffer` f64 values.
+ pub fn new(pool_size: usize, elements_per_buffer: usize) -> Result<Arc<Self>> {
+ if pool_size == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool requires at least one buffer".to_string(),
+ ));
+ }
+ if elements_per_buffer == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool buffer size must be greater than zero".to_string(),
+ ));
+ }
+
+ let mut buffers = Vec::with_capacity(pool_size);
+ for _ in 0..pool_size {
+ buffers.push(PinnedHostBuffer::new(elements_per_buffer)?);
+ }
+
+ Ok(Arc::new(Self {
+ free: Mutex::new(buffers),
+ available_cv: Condvar::new(),
+ capacity: pool_size,
+ elements_per_buffer,
+ }))
+ }
+
+ /// Acquire a pinned buffer, blocking until one is available.
+ pub fn acquire(self: &Arc<Self>) -> PinnedBufferHandle {
+ let mut free = self.free.lock().unwrap();
+ loop {
+ if let Some(buffer) = free.pop() {
+ return PinnedBufferHandle {
+ buffer: Some(buffer),
+ pool: Arc::clone(self),
+ };
+ }
+ free = self.available_cv.wait(free).unwrap();
+ }
+ }
+
+ /// Try to acquire a pinned buffer from the pool.
+ ///
+ /// Returns `None` if the pool is currently empty; callers can choose to spin/wait
+ /// or fall back to synchronous paths.
+ pub fn try_acquire(self: &Arc<Self>) -> Option<PinnedBufferHandle> {
+ let mut free = self.free.lock().unwrap();
+ free.pop().map(|buffer| PinnedBufferHandle {
+ buffer: Some(buffer),
+ pool: Arc::clone(self),
+ })
+ }
+
+ /// Number of buffers currently available.
+ pub fn available(&self) -> usize {
+ self.free.lock().unwrap().len()
Review Comment:
The `.unwrap()` call on the mutex lock operation can cause a panic if the mutex is poisoned. Consider handling this error more gracefully or documenting the panic behavior.
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -0,0 +1,141 @@
+//
+// 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.
+
+//! Reusable pool of pinned host buffers for staging Disk → Host → GPU transfers.
+//! Intended for producer/consumer pipelines that need a small, fixed set of
+//! page-locked buffers to avoid repeated cudaHostAlloc / cudaFreeHost.
+
+use std::sync::{Arc, Condvar, Mutex};
+
+use crate::error::{MahoutError, Result};
+use crate::gpu::memory::PinnedHostBuffer;
+
+/// Handle that automatically returns a buffer to the pool on drop.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferHandle {
+ buffer: Option<PinnedHostBuffer>,
+ pool: Arc<PinnedBufferPool>,
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::Deref for PinnedBufferHandle {
+ type Target = PinnedHostBuffer;
+
+ fn deref(&self) -> &Self::Target {
+ self.buffer
+ .as_ref()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::DerefMut for PinnedBufferHandle {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ self.buffer
+ .as_mut()
+ .expect("Buffer already returned to pool")
Review Comment:
The panic message "Buffer already returned to pool" may be misleading. This panic occurs when attempting to use a `PinnedBufferHandle` after it has been dropped and its buffer returned to the pool. Consider a more descriptive message such as "Attempted to use PinnedBufferHandle after buffer was returned to pool (use-after-drop)" to better indicate the programmer error.
##########
qdp/qdp-core/src/lib.rs:
##########
@@ -343,10 +327,16 @@ impl QdpEngine {
.cast::<u8>()
.add(offset_bytes)
.cast::<std::ffi::c_void>();
- debug_assert!(
- samples_in_chunk <= max_samples_per_chunk,
- "samples_in_chunk must be <= max_samples_per_chunk"
- );
+
+ let mut norm_buffer = self
+ .device
+ .alloc_zeros::<f64>(samples_in_chunk)
+ .map_err(|e| {
+ MahoutError::MemoryAllocation(format!(
+ "Failed to allocate norm buffer: {:?}",
+ e
+ ))
+ })?;
Review Comment:
Moving the norm buffer allocation inside the per-chunk loop (line 331-339) reintroduces per-chunk allocation overhead that was previously avoided. The old code comment explicitly stated: "Reuse a single norm buffer across chunks to avoid per-chunk allocations" and warned that "per-chunk allocation + drop can lead to use-after-free when the next chunk reuses the same device memory while the previous chunk is still running." While the use-after-free concern is mitigated by proper stream synchronization, the per-chunk allocation overhead remains. Consider pre-allocating a single norm buffer sized for the maximum expected samples_in_chunk to improve performance.
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -0,0 +1,141 @@
+//
+// 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.
+
+//! Reusable pool of pinned host buffers for staging Disk → Host → GPU transfers.
+//! Intended for producer/consumer pipelines that need a small, fixed set of
+//! page-locked buffers to avoid repeated cudaHostAlloc / cudaFreeHost.
+
+use std::sync::{Arc, Condvar, Mutex};
+
+use crate::error::{MahoutError, Result};
+use crate::gpu::memory::PinnedHostBuffer;
+
+/// Handle that automatically returns a buffer to the pool on drop.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferHandle {
+ buffer: Option<PinnedHostBuffer>,
+ pool: Arc<PinnedBufferPool>,
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::Deref for PinnedBufferHandle {
+ type Target = PinnedHostBuffer;
+
+ fn deref(&self) -> &Self::Target {
+ self.buffer
+ .as_ref()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::DerefMut for PinnedBufferHandle {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ self.buffer
+ .as_mut()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl Drop for PinnedBufferHandle {
+ fn drop(&mut self) {
+ if let Some(buf) = self.buffer.take() {
+ let mut free = self.pool.free.lock().unwrap();
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
+ }
+}
+
+/// Pool of pinned host buffers sized for a fixed batch shape.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferPool {
+ free: Mutex<Vec<PinnedHostBuffer>>,
+ available_cv: Condvar,
+ capacity: usize,
+ elements_per_buffer: usize,
+}
+
+#[cfg(target_os = "linux")]
+impl PinnedBufferPool {
+ /// Create a pool with `pool_size` pinned buffers, each sized for `elements_per_buffer` f64 values.
+ pub fn new(pool_size: usize, elements_per_buffer: usize) -> Result<Arc<Self>> {
+ if pool_size == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool requires at least one buffer".to_string(),
+ ));
+ }
+ if elements_per_buffer == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool buffer size must be greater than zero".to_string(),
+ ));
+ }
+
+ let mut buffers = Vec::with_capacity(pool_size);
+ for _ in 0..pool_size {
+ buffers.push(PinnedHostBuffer::new(elements_per_buffer)?);
+ }
+
+ Ok(Arc::new(Self {
+ free: Mutex::new(buffers),
+ available_cv: Condvar::new(),
+ capacity: pool_size,
+ elements_per_buffer,
+ }))
+ }
+
+ /// Acquire a pinned buffer, blocking until one is available.
+ pub fn acquire(self: &Arc<Self>) -> PinnedBufferHandle {
+ let mut free = self.free.lock().unwrap();
+ loop {
+ if let Some(buffer) = free.pop() {
+ return PinnedBufferHandle {
+ buffer: Some(buffer),
+ pool: Arc::clone(self),
+ };
+ }
+ free = self.available_cv.wait(free).unwrap();
+ }
+ }
+
+ /// Try to acquire a pinned buffer from the pool.
+ ///
+ /// Returns `None` if the pool is currently empty; callers can choose to spin/wait
+ /// or fall back to synchronous paths.
+ pub fn try_acquire(self: &Arc<Self>) -> Option<PinnedBufferHandle> {
+ let mut free = self.free.lock().unwrap();
Review Comment:
The `.unwrap()` call on the mutex lock operation can cause a panic if the mutex is poisoned. Consider handling this error more gracefully or documenting the panic behavior.
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -0,0 +1,141 @@
+//
+// 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.
+
+//! Reusable pool of pinned host buffers for staging Disk → Host → GPU transfers.
+//! Intended for producer/consumer pipelines that need a small, fixed set of
+//! page-locked buffers to avoid repeated cudaHostAlloc / cudaFreeHost.
+
+use std::sync::{Arc, Condvar, Mutex};
+
+use crate::error::{MahoutError, Result};
+use crate::gpu::memory::PinnedHostBuffer;
+
+/// Handle that automatically returns a buffer to the pool on drop.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferHandle {
+ buffer: Option<PinnedHostBuffer>,
+ pool: Arc<PinnedBufferPool>,
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::Deref for PinnedBufferHandle {
+ type Target = PinnedHostBuffer;
+
+ fn deref(&self) -> &Self::Target {
+ self.buffer
+ .as_ref()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::DerefMut for PinnedBufferHandle {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ self.buffer
+ .as_mut()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl Drop for PinnedBufferHandle {
+ fn drop(&mut self) {
+ if let Some(buf) = self.buffer.take() {
+ let mut free = self.pool.free.lock().unwrap();
Review Comment:
The `.unwrap()` call on the mutex lock operation can cause a panic if the mutex is poisoned. Consider handling this error more gracefully or documenting the panic behavior.
```suggestion
let mut free = match self.pool.free.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
```
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -0,0 +1,141 @@
+//
+// 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.
+
+//! Reusable pool of pinned host buffers for staging Disk → Host → GPU transfers.
+//! Intended for producer/consumer pipelines that need a small, fixed set of
+//! page-locked buffers to avoid repeated cudaHostAlloc / cudaFreeHost.
+
+use std::sync::{Arc, Condvar, Mutex};
+
+use crate::error::{MahoutError, Result};
+use crate::gpu::memory::PinnedHostBuffer;
+
+/// Handle that automatically returns a buffer to the pool on drop.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferHandle {
+ buffer: Option<PinnedHostBuffer>,
+ pool: Arc<PinnedBufferPool>,
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::Deref for PinnedBufferHandle {
+ type Target = PinnedHostBuffer;
+
+ fn deref(&self) -> &Self::Target {
+ self.buffer
+ .as_ref()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::DerefMut for PinnedBufferHandle {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ self.buffer
+ .as_mut()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl Drop for PinnedBufferHandle {
+ fn drop(&mut self) {
+ if let Some(buf) = self.buffer.take() {
+ let mut free = self.pool.free.lock().unwrap();
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
+ }
+}
+
+/// Pool of pinned host buffers sized for a fixed batch shape.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferPool {
+ free: Mutex<Vec<PinnedHostBuffer>>,
+ available_cv: Condvar,
+ capacity: usize,
+ elements_per_buffer: usize,
+}
+
+#[cfg(target_os = "linux")]
+impl PinnedBufferPool {
+ /// Create a pool with `pool_size` pinned buffers, each sized for `elements_per_buffer` f64 values.
+ pub fn new(pool_size: usize, elements_per_buffer: usize) -> Result<Arc<Self>> {
+ if pool_size == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool requires at least one buffer".to_string(),
+ ));
+ }
+ if elements_per_buffer == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool buffer size must be greater than zero".to_string(),
+ ));
+ }
+
+ let mut buffers = Vec::with_capacity(pool_size);
+ for _ in 0..pool_size {
+ buffers.push(PinnedHostBuffer::new(elements_per_buffer)?);
+ }
+
+ Ok(Arc::new(Self {
+ free: Mutex::new(buffers),
+ available_cv: Condvar::new(),
+ capacity: pool_size,
+ elements_per_buffer,
+ }))
+ }
+
+ /// Acquire a pinned buffer, blocking until one is available.
+ pub fn acquire(self: &Arc<Self>) -> PinnedBufferHandle {
+ let mut free = self.free.lock().unwrap();
+ loop {
+ if let Some(buffer) = free.pop() {
+ return PinnedBufferHandle {
+ buffer: Some(buffer),
+ pool: Arc::clone(self),
+ };
+ }
+ free = self.available_cv.wait(free).unwrap();
Review Comment:
The `.unwrap()` calls on mutex lock operations can cause panics if the mutex is poisoned. In a production system, poisoned mutex errors should be handled more gracefully, either by propagating the error or by documenting that panic behavior is intentional in these scenarios.
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -0,0 +1,141 @@
+//
+// 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.
+
+//! Reusable pool of pinned host buffers for staging Disk → Host → GPU transfers.
+//! Intended for producer/consumer pipelines that need a small, fixed set of
+//! page-locked buffers to avoid repeated cudaHostAlloc / cudaFreeHost.
+
+use std::sync::{Arc, Condvar, Mutex};
+
+use crate::error::{MahoutError, Result};
+use crate::gpu::memory::PinnedHostBuffer;
+
+/// Handle that automatically returns a buffer to the pool on drop.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferHandle {
+ buffer: Option<PinnedHostBuffer>,
+ pool: Arc<PinnedBufferPool>,
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::Deref for PinnedBufferHandle {
+ type Target = PinnedHostBuffer;
+
+ fn deref(&self) -> &Self::Target {
+ self.buffer
+ .as_ref()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl std::ops::DerefMut for PinnedBufferHandle {
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ self.buffer
+ .as_mut()
+ .expect("Buffer already returned to pool")
+ }
+}
+
+#[cfg(target_os = "linux")]
+impl Drop for PinnedBufferHandle {
+ fn drop(&mut self) {
+ if let Some(buf) = self.buffer.take() {
+ let mut free = self.pool.free.lock().unwrap();
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
+ }
+}
+
+/// Pool of pinned host buffers sized for a fixed batch shape.
+#[cfg(target_os = "linux")]
+pub struct PinnedBufferPool {
+ free: Mutex<Vec<PinnedHostBuffer>>,
+ available_cv: Condvar,
+ capacity: usize,
+ elements_per_buffer: usize,
+}
+
+#[cfg(target_os = "linux")]
+impl PinnedBufferPool {
+ /// Create a pool with `pool_size` pinned buffers, each sized for `elements_per_buffer` f64 values.
+ pub fn new(pool_size: usize, elements_per_buffer: usize) -> Result<Arc<Self>> {
+ if pool_size == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool requires at least one buffer".to_string(),
+ ));
+ }
+ if elements_per_buffer == 0 {
+ return Err(MahoutError::InvalidInput(
+ "PinnedBufferPool buffer size must be greater than zero".to_string(),
+ ));
+ }
+
+ let mut buffers = Vec::with_capacity(pool_size);
+ for _ in 0..pool_size {
+ buffers.push(PinnedHostBuffer::new(elements_per_buffer)?);
+ }
+
+ Ok(Arc::new(Self {
+ free: Mutex::new(buffers),
+ available_cv: Condvar::new(),
+ capacity: pool_size,
+ elements_per_buffer,
+ }))
+ }
+
+ /// Acquire a pinned buffer, blocking until one is available.
+ pub fn acquire(self: &Arc<Self>) -> PinnedBufferHandle {
+ let mut free = self.free.lock().unwrap();
+ loop {
+ if let Some(buffer) = free.pop() {
+ return PinnedBufferHandle {
+ buffer: Some(buffer),
+ pool: Arc::clone(self),
+ };
+ }
+ free = self.available_cv.wait(free).unwrap();
+ }
+ }
+
+ /// Try to acquire a pinned buffer from the pool.
+ ///
+ /// Returns `None` if the pool is currently empty; callers can choose to spin/wait
+ /// or fall back to synchronous paths.
+ pub fn try_acquire(self: &Arc<Self>) -> Option<PinnedBufferHandle> {
+ let mut free = self.free.lock().unwrap();
+ free.pop().map(|buffer| PinnedBufferHandle {
+ buffer: Some(buffer),
+ pool: Arc::clone(self),
+ })
+ }
+
+ /// Number of buffers currently available.
+ pub fn available(&self) -> usize {
+ self.free.lock().unwrap().len()
+ }
+
+ /// Total number of buffers managed by this pool.
+ pub fn capacity(&self) -> usize {
+ self.capacity
+ }
+
+ /// Fixed element capacity for each buffer in the pool.
+ pub fn elements_per_buffer(&self) -> usize {
+ self.elements_per_buffer
+ }
+}
Review Comment:
The new `PinnedBufferPool` implementation lacks test coverage. Consider adding tests to verify buffer acquisition/release patterns, concurrent access behavior, pool exhaustion scenarios, and error handling (e.g., zero pool size, zero buffer size).
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3698318462
plz fix pre-commit error
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3707062929
nits: in `qdp/qdp-core/src/lib.rs` 38-41 lines seem not to be used and could be removed.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "ryankert01 (via GitHub)" <gi...@apache.org>.
ryankert01 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2657598368
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -54,9 +54,11 @@ impl std::ops::DerefMut for PinnedBufferHandle {
impl Drop for PinnedBufferHandle {
fn drop(&mut self) {
if let Some(buf) = self.buffer.take() {
- let mut free = self.pool.free.lock().unwrap();
- free.push(buf);
- self.pool.available_cv.notify_one();
+ // If the mutex is poisoned, avoid touching shared state during drop.
+ if let Ok(mut free) = self.pool.free.lock() {
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
Review Comment:
If the mutex is poisoned, we can get the lock anyways during drop? Like `lock_free()` defined in L103.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3693171582
I think maybe we could add some unit tests for this.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2650986384
##########
qdp/qdp-core/src/gpu/pipeline.rs:
##########
@@ -24,90 +24,116 @@ use std::ffi::c_void;
use cudarc::driver::{CudaDevice, CudaSlice, DevicePtr, safe::CudaStream};
use crate::error::{MahoutError, Result};
#[cfg(target_os = "linux")]
-use crate::gpu::memory::{ensure_device_memory_available, map_allocation_error, PinnedBuffer};
-
+use crate::gpu::memory::{ensure_device_memory_available, map_allocation_error};
+#[cfg(target_os = "linux")]
+use crate::gpu::buffer_pool::{PinnedBufferPool, PinnedBufferHandle};
#[cfg(target_os = "linux")]
use crate::gpu::cuda_ffi::{
- cudaEventCreateWithFlags, cudaEventDestroy, cudaEventRecord, cudaMemcpyAsync, cudaStreamSynchronize,
- cudaStreamWaitEvent, CUDA_EVENT_DISABLE_TIMING, CUDA_MEMCPY_HOST_TO_DEVICE,
+ cudaEventCreateWithFlags,
+ cudaEventDestroy,
+ cudaEventRecord,
+ cudaMemcpyAsync,
+ cudaStreamSynchronize,
+ cudaStreamWaitEvent,
+ CUDA_EVENT_DISABLE_TIMING,
+ CUDA_MEMCPY_HOST_TO_DEVICE,
};
-/// Dual-stream pipeline context: manages compute/copy streams and sync events
+/// Dual-stream context coordinating copy/compute with an event.
#[cfg(target_os = "linux")]
pub struct PipelineContext {
pub stream_compute: CudaStream,
pub stream_copy: CudaStream,
- event_copy_done: *mut c_void,
+ events_copy_done: Vec<*mut c_void>,
}
#[cfg(target_os = "linux")]
+#[allow(unsafe_op_in_unsafe_fn)]
impl PipelineContext {
- /// Create dual streams and sync event
- pub fn new(device: &Arc<CudaDevice>) -> Result<Self> {
+ pub fn new(device: &Arc<CudaDevice>, event_slots: usize) -> Result<Self> {
let stream_compute = device.fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("{:?}", e)))?;
+ .map_err(|e| MahoutError::Cuda(format!("Failed to create compute stream: {:?}", e)))?;
let stream_copy = device.fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("{:?}", e)))?;
+ .map_err(|e| MahoutError::Cuda(format!("Failed to create copy stream: {:?}", e)))?;
- let mut event_copy_done: *mut c_void = std::ptr::null_mut();
- unsafe {
- let ret = cudaEventCreateWithFlags(&mut event_copy_done, CUDA_EVENT_DISABLE_TIMING);
- if ret != 0 {
- return Err(MahoutError::Cuda(format!("Failed to create CUDA event: {}", ret)));
+ let mut events_copy_done = Vec::with_capacity(event_slots);
+ for _ in 0..event_slots {
+ let mut ev: *mut c_void = std::ptr::null_mut();
+ unsafe {
+ let ret = cudaEventCreateWithFlags(&mut ev, CUDA_EVENT_DISABLE_TIMING);
+ if ret != 0 {
+ return Err(MahoutError::Cuda(format!("Failed to create CUDA event: {}", ret)));
+ }
}
+ events_copy_done.push(ev);
}
Ok(Self {
stream_compute,
stream_copy,
- event_copy_done,
+ events_copy_done,
})
}
- /// Async H2D copy on copy stream
- pub unsafe fn async_copy_to_device(&self, src: &PinnedBuffer, dst: *mut c_void, len_elements: usize) {
+ /// Async H2D copy on the copy stream.
+ pub unsafe fn async_copy_to_device(
+ &self,
+ src: *const c_void,
+ dst: *mut c_void,
+ len_elements: usize,
+ ) -> Result<()> {
crate::profile_scope!("GPU::H2D_Copy");
- unsafe {
- cudaMemcpyAsync(
- dst,
- src.ptr() as *const c_void,
- len_elements * std::mem::size_of::<f64>(),
- CUDA_MEMCPY_HOST_TO_DEVICE,
- self.stream_copy.stream as *mut c_void
- );
+ let ret = cudaMemcpyAsync(
+ dst,
+ src,
+ len_elements * std::mem::size_of::<f64>(),
+ CUDA_MEMCPY_HOST_TO_DEVICE,
+ self.stream_copy.stream as *mut c_void,
+ );
+ if ret != 0 {
+ return Err(MahoutError::Cuda(format!("Async H2D copy failed with CUDA error: {}", ret)));
}
+ Ok(())
}
- /// Record copy completion event
- pub unsafe fn record_copy_done(&self) {
- unsafe {
- cudaEventRecord(self.event_copy_done, self.stream_copy.stream as *mut c_void);
+ /// Record completion of the copy on the copy stream.
+ pub unsafe fn record_copy_done(&self, slot: usize) -> Result<()> {
+ let ret = cudaEventRecord(self.events_copy_done[slot], self.stream_copy.stream as *mut c_void);
+ if ret != 0 {
+ return Err(MahoutError::Cuda(format!("cudaEventRecord failed: {}", ret)));
}
+ Ok(())
}
- /// Make compute stream wait for copy completion
- pub unsafe fn wait_for_copy(&self) {
+ /// Make compute stream wait for the copy completion event.
+ pub unsafe fn wait_for_copy(&self, slot: usize) -> Result<()> {
crate::profile_scope!("GPU::StreamWait");
- unsafe {
- cudaStreamWaitEvent(self.stream_compute.stream as *mut c_void, self.event_copy_done, 0);
+ let ret = cudaStreamWaitEvent(self.stream_compute.stream as *mut c_void, self.events_copy_done[slot], 0);
Review Comment:
done
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2658926866
##########
qdp/qdp-core/src/lib.rs:
##########
@@ -221,38 +228,13 @@ impl QdpEngine {
"Sample size cannot be zero".into(),
));
}
- if sample_size > STAGE_SIZE_ELEMENTS {
- return Err(MahoutError::InvalidInput(format!(
- "Sample size {} exceeds staging buffer capacity {} (elements)",
- sample_size, STAGE_SIZE_ELEMENTS
- )));
- }
-
- // Reuse a single norm buffer across chunks to avoid per-chunk allocations.
- //
- // Important: the norm buffer must outlive the async kernels that consume it.
- // Per-chunk allocation + drop can lead to use-after-free when the next chunk
- // reuses the same device memory while the previous chunk is still running.
Review Comment:
I think here has a potential problem in `qdp/qdp-core/src/lib.rs`. In the `encode_from_parquet()` function in `qdp/qdp-core/src/lib.rs`, there is a critical use-after-free bug in the lifetime management of `norm_buffer`. The code allocates `norm_buffer` inside the `BatchEncode` scope at line 331-339, and asynchronously launches `launch_l2_norm_batch` and `launch_amplitude_encode_batch` kernels at line 343-375, which execute asynchronously on `ctx.stream_compute`. However, when the `BatchEncode` scope ends at line 376, `norm_buffer` is immediately dropped, and according to the comment in `pipeline.rs:336`, when `CudaSlice` is dropped, it immediately calls `cudaFree` to free GPU memory. The problem is that `ctx.sync_copy_stream()` at line 378 only synchronizes the copy stream, not the compute stream, so when `norm_buffer` is freed, kernels on the compute stream may still be executing, causing kernels to access GPU memory that has already been freed. Even though kernels execute seque
ntially on the same stream, if their execution time is long, `norm_buffer` may still be freed before the kernels complete. In a loop processing multiple chunks, the first chunk's `norm_buffer` is dropped before its kernel completes, and while the second chunk's kernel will wait for the first to complete, the first kernel is already accessing freed memory. The old code's comment explicitly warned about this issue, stating that the norm buffer must outlive the async kernels that consume it, and that per-chunk allocation plus drop can lead to use-after-free. Plz correct me if I'm wrong.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Pinned host buffer + dual-stream event pipeline to overlap copy and compute [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3684913542
Thanks @400Ping for the patch!
1. What's the reason you define a ffi function again?
2. some tests failed locally due to tensor shape problem
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "ryankert01 (via GitHub)" <gi...@apache.org>.
ryankert01 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3703358375
please fix pre-commit. I tested locally and get a 2.8% speedup on arrow ipc case.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "guan404ming (via GitHub)" <gi...@apache.org>.
guan404ming commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3703868137
I agree that the comment regarding `.unwrap()` is valid.
Do we want to handle this more gracefully or is the current panic-on-poison behavior expected? If so, we can document it.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "ryankert01 (via GitHub)" <gi...@apache.org>.
ryankert01 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2657588261
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -54,9 +54,11 @@ impl std::ops::DerefMut for PinnedBufferHandle {
impl Drop for PinnedBufferHandle {
fn drop(&mut self) {
if let Some(buf) = self.buffer.take() {
- let mut free = self.pool.free.lock().unwrap();
- free.push(buf);
- self.pool.available_cv.notify_one();
+ // If the mutex is poisoned, avoid touching shared state during drop.
+ if let Ok(mut free) = self.pool.free.lock() {
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
Review Comment:
```diff
let mut free = self.pool.free
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
free.push(buf);
self.pool.available_cv.notify_one();
```
##########
qdp/qdp-core/src/gpu/buffer_pool.rs:
##########
@@ -54,9 +54,11 @@ impl std::ops::DerefMut for PinnedBufferHandle {
impl Drop for PinnedBufferHandle {
fn drop(&mut self) {
if let Some(buf) = self.buffer.take() {
- let mut free = self.pool.free.lock().unwrap();
- free.push(buf);
- self.pool.available_cv.notify_one();
+ // If the mutex is poisoned, avoid touching shared state during drop.
+ if let Ok(mut free) = self.pool.free.lock() {
+ free.push(buf);
+ self.pool.available_cv.notify_one();
+ }
Review Comment:
```diff
let mut free = self.pool.free
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
free.push(buf);
self.pool.available_cv.notify_one();
```
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3707233144
> thanks for the update! There are many return Err() functions in this PR. but we don't actually sync compute stream. I think we should prevent this in other PR.
Should I do this in this pr or open up a follow up
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Pinned host buffer + dual-stream event pipeline to overlap copy and compute [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3686079525
> Thanks @400Ping for the patch!
>
> 1. What's the reason you define a ffi function again?
> 2. some tests failed locally due to tensor shape problem or you should fix test for excepted output.
My bad just fixed it.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2648437169
##########
qdp/qdp-core/src/gpu/pipeline.rs:
##########
@@ -24,90 +24,116 @@ use std::ffi::c_void;
use cudarc::driver::{CudaDevice, CudaSlice, DevicePtr, safe::CudaStream};
use crate::error::{MahoutError, Result};
#[cfg(target_os = "linux")]
-use crate::gpu::memory::{ensure_device_memory_available, map_allocation_error, PinnedBuffer};
-
+use crate::gpu::memory::{ensure_device_memory_available, map_allocation_error};
+#[cfg(target_os = "linux")]
+use crate::gpu::buffer_pool::{PinnedBufferPool, PinnedBufferHandle};
#[cfg(target_os = "linux")]
use crate::gpu::cuda_ffi::{
- cudaEventCreateWithFlags, cudaEventDestroy, cudaEventRecord, cudaMemcpyAsync, cudaStreamSynchronize,
- cudaStreamWaitEvent, CUDA_EVENT_DISABLE_TIMING, CUDA_MEMCPY_HOST_TO_DEVICE,
+ cudaEventCreateWithFlags,
+ cudaEventDestroy,
+ cudaEventRecord,
+ cudaMemcpyAsync,
+ cudaStreamSynchronize,
+ cudaStreamWaitEvent,
+ CUDA_EVENT_DISABLE_TIMING,
+ CUDA_MEMCPY_HOST_TO_DEVICE,
};
-/// Dual-stream pipeline context: manages compute/copy streams and sync events
+/// Dual-stream context coordinating copy/compute with an event.
#[cfg(target_os = "linux")]
pub struct PipelineContext {
pub stream_compute: CudaStream,
pub stream_copy: CudaStream,
- event_copy_done: *mut c_void,
+ events_copy_done: Vec<*mut c_void>,
}
#[cfg(target_os = "linux")]
+#[allow(unsafe_op_in_unsafe_fn)]
impl PipelineContext {
- /// Create dual streams and sync event
- pub fn new(device: &Arc<CudaDevice>) -> Result<Self> {
+ pub fn new(device: &Arc<CudaDevice>, event_slots: usize) -> Result<Self> {
let stream_compute = device.fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("{:?}", e)))?;
+ .map_err(|e| MahoutError::Cuda(format!("Failed to create compute stream: {:?}", e)))?;
let stream_copy = device.fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("{:?}", e)))?;
+ .map_err(|e| MahoutError::Cuda(format!("Failed to create copy stream: {:?}", e)))?;
- let mut event_copy_done: *mut c_void = std::ptr::null_mut();
- unsafe {
- let ret = cudaEventCreateWithFlags(&mut event_copy_done, CUDA_EVENT_DISABLE_TIMING);
- if ret != 0 {
- return Err(MahoutError::Cuda(format!("Failed to create CUDA event: {}", ret)));
+ let mut events_copy_done = Vec::with_capacity(event_slots);
+ for _ in 0..event_slots {
+ let mut ev: *mut c_void = std::ptr::null_mut();
+ unsafe {
+ let ret = cudaEventCreateWithFlags(&mut ev, CUDA_EVENT_DISABLE_TIMING);
+ if ret != 0 {
+ return Err(MahoutError::Cuda(format!("Failed to create CUDA event: {}", ret)));
+ }
}
+ events_copy_done.push(ev);
}
Ok(Self {
stream_compute,
stream_copy,
- event_copy_done,
+ events_copy_done,
})
}
- /// Async H2D copy on copy stream
- pub unsafe fn async_copy_to_device(&self, src: &PinnedBuffer, dst: *mut c_void, len_elements: usize) {
+ /// Async H2D copy on the copy stream.
+ pub unsafe fn async_copy_to_device(
+ &self,
+ src: *const c_void,
+ dst: *mut c_void,
+ len_elements: usize,
+ ) -> Result<()> {
crate::profile_scope!("GPU::H2D_Copy");
- unsafe {
- cudaMemcpyAsync(
- dst,
- src.ptr() as *const c_void,
- len_elements * std::mem::size_of::<f64>(),
- CUDA_MEMCPY_HOST_TO_DEVICE,
- self.stream_copy.stream as *mut c_void
- );
+ let ret = cudaMemcpyAsync(
+ dst,
+ src,
+ len_elements * std::mem::size_of::<f64>(),
+ CUDA_MEMCPY_HOST_TO_DEVICE,
+ self.stream_copy.stream as *mut c_void,
+ );
+ if ret != 0 {
+ return Err(MahoutError::Cuda(format!("Async H2D copy failed with CUDA error: {}", ret)));
}
+ Ok(())
}
- /// Record copy completion event
- pub unsafe fn record_copy_done(&self) {
- unsafe {
- cudaEventRecord(self.event_copy_done, self.stream_copy.stream as *mut c_void);
+ /// Record completion of the copy on the copy stream.
+ pub unsafe fn record_copy_done(&self, slot: usize) -> Result<()> {
+ let ret = cudaEventRecord(self.events_copy_done[slot], self.stream_copy.stream as *mut c_void);
+ if ret != 0 {
+ return Err(MahoutError::Cuda(format!("cudaEventRecord failed: {}", ret)));
}
+ Ok(())
}
- /// Make compute stream wait for copy completion
- pub unsafe fn wait_for_copy(&self) {
+ /// Make compute stream wait for the copy completion event.
+ pub unsafe fn wait_for_copy(&self, slot: usize) -> Result<()> {
crate::profile_scope!("GPU::StreamWait");
- unsafe {
- cudaStreamWaitEvent(self.stream_compute.stream as *mut c_void, self.event_copy_done, 0);
+ let ret = cudaStreamWaitEvent(self.stream_compute.stream as *mut c_void, self.events_copy_done[slot], 0);
Review Comment:
should we add a check here?
like:
```
if slot >= self.events_copy_done.len() {
return Err(MahoutError::InvalidInput(format!(
"Event slot {} out of range (max: {})",
slot, self.events_copy_done.len() - 1
)));
}
```
or something
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered async I/O for read_parquet_batch [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3682166772
cc @guan404ming @rich7420 @ryankert01
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2658918782
##########
qdp/qdp-core/src/gpu/pipeline.rs:
##########
@@ -191,32 +252,40 @@ where
{
crate::profile_scope!("GPU::AsyncPipeline");
- // 1. Create dual streams for pipeline overlap
- let stream1 = device
- .fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("Failed to create stream 1: {:?}", e)))?;
- let stream2 = device
- .fork_default_stream()
- .map_err(|e| MahoutError::Cuda(format!("Failed to create stream 2: {:?}", e)))?;
- let streams = [&stream1, &stream2];
+ // Pinned host staging pool sized to the current chunking strategy (double-buffer by default).
+ const CHUNK_SIZE_ELEMENTS: usize = 8 * 1024 * 1024 / std::mem::size_of::<f64>(); // 8MB
+ const PINNED_POOL_SIZE: usize = 2; // double buffering
+ // 1. Create dual streams with per-slot events to coordinate copy -> compute
+ let ctx = PipelineContext::new(device, PINNED_POOL_SIZE)?;
+ let pinned_pool = PinnedBufferPool::new(PINNED_POOL_SIZE, CHUNK_SIZE_ELEMENTS)
+ .map_err(|e| MahoutError::Cuda(format!("Failed to create pinned buffer pool: {}", e)))?;
// 2. Chunk size: 8MB per chunk (balance between overhead and overlap opportunity)
- // TODO: we should tune this dynamically based on the detected GPU model or PCIe bandwidth in the future.
- // Too small = launch overhead dominates, too large = less overlap
- const CHUNK_SIZE_ELEMENTS: usize = 8 * 1024 * 1024 / std::mem::size_of::<f64>(); // 8MB
+ // TODO: tune dynamically based on GPU/PCIe bandwidth.
// 3. Keep temporary buffers alive until all streams complete
// This prevents Rust from dropping them while GPU is still using them
let mut keep_alive_buffers: Vec<CudaSlice<f64>> = Vec::new();
+ // Keep pinned buffers alive until the copy stream has completed their H2D copy
+ let mut in_flight_pinned: Vec<PinnedBufferHandle> = Vec::new();
let mut global_offset = 0;
- // 4. Pipeline loop: alternate between streams for maximum overlap
+ // 4. Pipeline loop: copy on copy stream, compute on compute stream with event handoff
for (chunk_idx, chunk) in host_data.chunks(CHUNK_SIZE_ELEMENTS).enumerate() {
- let current_stream = streams[chunk_idx % 2];
+ let chunk_offset = global_offset;
+ let event_slot = chunk_idx % PINNED_POOL_SIZE;
crate::profile_scope!("GPU::ChunkProcess");
+ if chunk.len() > CHUNK_SIZE_ELEMENTS {
Review Comment:
I think this function would be never used. should be removed
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2658956834
##########
qdp/qdp-core/src/lib.rs:
##########
@@ -221,38 +228,13 @@ impl QdpEngine {
"Sample size cannot be zero".into(),
));
}
- if sample_size > STAGE_SIZE_ELEMENTS {
- return Err(MahoutError::InvalidInput(format!(
- "Sample size {} exceeds staging buffer capacity {} (elements)",
- sample_size, STAGE_SIZE_ELEMENTS
- )));
- }
-
- // Reuse a single norm buffer across chunks to avoid per-chunk allocations.
- //
- // Important: the norm buffer must outlive the async kernels that consume it.
- // Per-chunk allocation + drop can lead to use-after-free when the next chunk
- // reuses the same device memory while the previous chunk is still running.
Review Comment:
updated in [a99baf7](https://github.com/apache/mahout/pull/751/commits/a99baf7ea8992aad605e5ad3bfeec59ea7979d56)
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2658964691
##########
qdp/qdp-core/src/lib.rs:
##########
@@ -221,38 +228,13 @@ impl QdpEngine {
"Sample size cannot be zero".into(),
));
}
- if sample_size > STAGE_SIZE_ELEMENTS {
- return Err(MahoutError::InvalidInput(format!(
- "Sample size {} exceeds staging buffer capacity {} (elements)",
- sample_size, STAGE_SIZE_ELEMENTS
- )));
- }
-
- // Reuse a single norm buffer across chunks to avoid per-chunk allocations.
- //
- // Important: the norm buffer must outlive the async kernels that consume it.
- // Per-chunk allocation + drop can lead to use-after-free when the next chunk
- // reuses the same device memory while the previous chunk is still running.
Review Comment:
thanks for the update
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3708118732
Other pr
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "guan404ming (via GitHub)" <gi...@apache.org>.
guan404ming commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3708286437
Let's refine it after we get into main.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "guan404ming (via GitHub)" <gi...@apache.org>.
guan404ming commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3705057321
Need resolve conflicts, and overall looks good to me!
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3707120816
except that overall LGTM
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3703877698
> I agree that the comment regarding `.unwrap()` is valid. Do we want to handle this more gracefully, or is the current panic-on-poison behavior expected? If so, we can document it.
I think I will change the code to handle it more gracefully and add some comments to document this behavior.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3696624315
Just tested, the second one doesn't improve much performance, going to remove it.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3696529732
> We have 2 improvement in this PR. Based on the benchmark result, I'm speculating if there's one of them are not contributing to the speed improvement. What's your experience?
I think both have improvements, for the second one is what @rich7420 and @guan404ming suggested to change a different decompression technique to improve its performance. But I think overall it is because of the first one improving the speed improvements.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "ryankert01 (via GitHub)" <gi...@apache.org>.
ryankert01 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3693528868
We have 2 improvement in this PR. Based on the benchmark result, I'm speculating if there's one of them are not contributing to the speed improvement. What's your experience?
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2658955167
##########
qdp/qdp-core/src/lib.rs:
##########
@@ -221,38 +228,13 @@ impl QdpEngine {
"Sample size cannot be zero".into(),
));
}
- if sample_size > STAGE_SIZE_ELEMENTS {
- return Err(MahoutError::InvalidInput(format!(
- "Sample size {} exceeds staging buffer capacity {} (elements)",
- sample_size, STAGE_SIZE_ELEMENTS
- )));
- }
-
- // Reuse a single norm buffer across chunks to avoid per-chunk allocations.
- //
- // Important: the norm buffer must outlive the async kernels that consume it.
- // Per-chunk allocation + drop can lead to use-after-free when the next chunk
- // reuses the same device memory while the previous chunk is still running.
Review Comment:
hmm, I agree that it is a concern. Will update this into the pr
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "400Ping (via GitHub)" <gi...@apache.org>.
400Ping commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2647067332
##########
qdp/qdp-core/src/gpu/pipeline.rs:
##########
@@ -194,39 +216,40 @@ where
e,
))?;
+ // Acquire pinned staging buffer and populate it with the current chunk
+ let mut pinned_buf = pinned_pool.acquire();
+ pinned_buf.as_slice_mut()[..chunk.len()].copy_from_slice(chunk);
Review Comment:
done.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Pinned host buffer + dual-stream event pipeline to overlap copy and compute [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on code in PR #751:
URL: https://github.com/apache/mahout/pull/751#discussion_r2646940770
##########
qdp/qdp-core/src/gpu/pipeline.rs:
##########
@@ -194,39 +216,40 @@ where
e,
))?;
+ // Acquire pinned staging buffer and populate it with the current chunk
+ let mut pinned_buf = pinned_pool.acquire();
+ pinned_buf.as_slice_mut()[..chunk.len()].copy_from_slice(chunk);
Review Comment:
maybe we could add a check here.
like: `chunk.len() > CHUNK_SIZE_ELEMENTS`
##########
qdp/qdp-core/src/gpu/memory.rs:
##########
@@ -134,6 +136,78 @@ pub(crate) fn map_allocation_error(
}
}
+/// Page-locked host buffer used to enable true async DMA for H2D copies.
+#[cfg(target_os = "linux")]
+pub struct PinnedBuffer<'a, T> {
Review Comment:
nits: It seems not been used in any other files.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "rich7420 (via GitHub)" <gi...@apache.org>.
rich7420 commented on PR #751:
URL: https://github.com/apache/mahout/pull/751#issuecomment-3707120376
thanks for the update!
There are many return Err() functions in this PR. but we don't actually sync compute stream.
I think we should prevent this in other PR.
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [QDP] Double-buffered pinned I/O pipeline and faster Parquet decode [mahout]
Posted by "guan404ming (via GitHub)" <gi...@apache.org>.
guan404ming merged PR #751:
URL: https://github.com/apache/mahout/pull/751
--
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: dev-unsubscribe@mahout.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org