From c37ed3ff08e2b8f9ac7924d27c1c501eba6ea673 Mon Sep 17 00:00:00 2001 From: Shiv Bhatia Date: Sun, 17 May 2026 10:52:06 +0100 Subject: [PATCH] Sketch out gpu physical execution --- Cargo.lock | 14 ++ Cargo.toml | 2 + datafusion/gpu/Cargo.toml | 47 +++++ datafusion/gpu/src/batch.rs | 121 +++++++++++ datafusion/gpu/src/boundary.rs | 271 +++++++++++++++++++++++++ datafusion/gpu/src/execution_plan.rs | 77 +++++++ datafusion/gpu/src/lib.rs | 60 ++++++ datafusion/gpu/src/memory_pool.rs | 231 +++++++++++++++++++++ datafusion/gpu/src/operators.rs | 290 +++++++++++++++++++++++++++ datafusion/gpu/src/optimizer.rs | 67 +++++++ datafusion/gpu/src/stream.rs | 37 ++++ 11 files changed, 1217 insertions(+) create mode 100644 datafusion/gpu/Cargo.toml create mode 100644 datafusion/gpu/src/batch.rs create mode 100644 datafusion/gpu/src/boundary.rs create mode 100644 datafusion/gpu/src/execution_plan.rs create mode 100644 datafusion/gpu/src/lib.rs create mode 100644 datafusion/gpu/src/memory_pool.rs create mode 100644 datafusion/gpu/src/operators.rs create mode 100644 datafusion/gpu/src/optimizer.rs create mode 100644 datafusion/gpu/src/stream.rs diff --git a/Cargo.lock b/Cargo.lock index d2ce889675b1d..0992973cebf66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2295,6 +2295,20 @@ dependencies = [ "datafusion-physical-expr-common", ] +[[package]] +name = "datafusion-gpu" +version = "53.1.0" +dependencies = [ + "arrow", + "datafusion-common", + "datafusion-execution", + "datafusion-physical-expr", + "datafusion-physical-optimizer", + "datafusion-physical-plan", + "futures", + "tokio", +] + [[package]] name = "datafusion-macros" version = "53.1.0" diff --git a/Cargo.toml b/Cargo.toml index 78c271d524fb8..f952e712f740d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ members = [ "datafusion/expr-common", "datafusion/execution", "datafusion/ffi", + "datafusion/gpu", "datafusion/functions", "datafusion/functions-aggregate", "datafusion/functions-aggregate-common", @@ -143,6 +144,7 @@ datafusion-functions-nested = { path = "datafusion/functions-nested", version = datafusion-functions-table = { path = "datafusion/functions-table", version = "53.1.0" } datafusion-functions-window = { path = "datafusion/functions-window", version = "53.1.0" } datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "53.1.0" } +datafusion-gpu = { path = "datafusion/gpu", version = "53.1.0" } datafusion-macros = { path = "datafusion/macros", version = "53.1.0" } datafusion-optimizer = { path = "datafusion/optimizer", version = "53.1.0", default-features = false } datafusion-physical-expr = { path = "datafusion/physical-expr", version = "53.1.0", default-features = false } diff --git a/datafusion/gpu/Cargo.toml b/datafusion/gpu/Cargo.toml new file mode 100644 index 0000000000000..70abdb3d691e2 --- /dev/null +++ b/datafusion/gpu/Cargo.toml @@ -0,0 +1,47 @@ +# 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. + +[package] +name = "datafusion-gpu" +description = "DataFusion GPU physical operators (experimental)" +keywords = ["datafusion", "gpu", "cuda"] +version = { workspace = true } +edition = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +authors = { workspace = true } +rust-version = { workspace = true } +publish = false + +# Note: add additional linter rules in lib.rs. +# Rust does not support workspace + new linter rules in subcrates yet +# https://github.com/rust-lang/cargo/issues/13157 +[lints] +workspace = true + +[dependencies] +arrow = { workspace = true } +datafusion-common = { workspace = true } +datafusion-execution = { workspace = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-optimizer = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true } diff --git a/datafusion/gpu/src/batch.rs b/datafusion/gpu/src/batch.rs new file mode 100644 index 0000000000000..a282ed4b230da --- /dev/null +++ b/datafusion/gpu/src/batch.rs @@ -0,0 +1,121 @@ +// 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. + +//! Device-resident column buffers and record batches. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, internal_err}; + +use crate::memory_pool::GpuMemoryReservation; + +/// A column buffer resident in GPU device memory. +/// +/// In this scaffold the underlying storage is a host-side placeholder; the +/// real implementation will replace [`Backing::HostPlaceholder`] with a CUDA +/// device allocation (e.g. `cudarc::driver::CudaSlice`). Operators must +/// treat the bytes as opaque — they are *not* a valid host pointer for kernel +/// inputs. +#[derive(Debug)] +pub struct GpuBuffer { + byte_len: usize, + #[expect(dead_code, reason = "read once the CUDA backend lands")] + backing: Backing, +} + +#[derive(Debug)] +enum Backing { + /// Host-resident placeholder. Held only so the scaffold compiles and + /// round-trips data through `GpuUploadExec`/`GpuDownloadExec` once those + /// land; never to be dereferenced by GPU code. + HostPlaceholder( + #[expect(dead_code, reason = "read once the CUDA backend lands")] Vec, + ), +} + +impl GpuBuffer { + pub fn from_host_placeholder(bytes: Vec) -> Self { + Self { + byte_len: bytes.len(), + backing: Backing::HostPlaceholder(bytes), + } + } + + pub fn byte_len(&self) -> usize { + self.byte_len + } +} + +/// Arrow-shaped batch whose column buffers live in device memory. +/// +/// This is the unit of data exchanged between operators implementing +/// [`crate::GpuExecutionPlan`]. It deliberately does *not* implement any +/// conversion to `arrow::record_batch::RecordBatch`; crossing the host/device +/// boundary must go through [`crate::boundary::GpuUploadExec`] / +/// [`crate::boundary::GpuDownloadExec`]. +#[derive(Debug)] +pub struct GpuRecordBatch { + schema: SchemaRef, + columns: Vec>, + num_rows: usize, + // Dropped with the batch, releasing the device-memory reservation. + _reservation: GpuMemoryReservation, +} + +impl GpuRecordBatch { + pub fn try_new( + schema: SchemaRef, + columns: Vec>, + num_rows: usize, + reservation: GpuMemoryReservation, + ) -> Result { + if schema.fields().len() != columns.len() { + return internal_err!( + "GpuRecordBatch: schema has {} fields but {} columns supplied", + schema.fields().len(), + columns.len() + ); + } + Ok(Self { + schema, + columns, + num_rows, + _reservation: reservation, + }) + } + + pub fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + pub fn num_rows(&self) -> usize { + self.num_rows + } + + pub fn num_columns(&self) -> usize { + self.columns.len() + } + + pub fn column(&self, idx: usize) -> &Arc { + &self.columns[idx] + } + + pub fn columns(&self) -> &[Arc] { + &self.columns + } +} diff --git a/datafusion/gpu/src/boundary.rs b/datafusion/gpu/src/boundary.rs new file mode 100644 index 0000000000000..a7c4c39e0e4c9 --- /dev/null +++ b/datafusion/gpu/src/boundary.rs @@ -0,0 +1,271 @@ +// 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. + +//! Boundary operators that move data across the host/device line. +//! +//! [`GpuUploadExec`] consumes a CPU [`ExecutionPlan`] child and produces a +//! stream of device-resident [`crate::GpuRecordBatch`]es. +//! [`GpuDownloadExec`] is the inverse — it consumes a +//! [`crate::GpuExecutionPlan`] child and produces a normal CPU +//! `RecordBatchStream`. +//! +//! These are the *only* places where host and device data legally meet. The +//! [`crate::optimizer::InsertGpuExec`] rule is responsible for inserting them +//! at the edges of every GPU subtree it creates. + +use std::any::Any; +use std::fmt; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, PlanProperties, +}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, +}; + +use crate::execution_plan::{GpuExecutionPlan, as_gpu_execution_plan}; +use crate::stream::SendableGpuRecordBatchStream; + +fn passthrough_plan_properties(schema: SchemaRef) -> Arc { + Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )) +} + +/// CPU → GPU boundary. Pulls host [`RecordBatch`]es from `child`, copies their +/// buffers to device memory, and emits [`crate::GpuRecordBatch`]es. +#[derive(Debug)] +pub struct GpuUploadExec { + child: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl GpuUploadExec { + pub fn new(child: Arc) -> Self { + let schema = child.schema(); + let properties = passthrough_plan_properties(Arc::clone(&schema)); + Self { + child, + schema, + properties, + } + } + + pub fn input(&self) -> &Arc { + &self.child + } +} + +impl DisplayAs for GpuUploadExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "GpuUploadExec") + } +} + +impl ExecutionPlan for GpuUploadExec { + fn name(&self) -> &str { + "GpuUploadExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("GpuUploadExec must have exactly one child"); + } + let [child] = children.try_into().unwrap(); + Ok(Arc::new(GpuUploadExec::new(child))) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + internal_err!( + "GpuUploadExec produces device-resident batches; call gpu_execute or \ + wrap in GpuDownloadExec before consuming on the host" + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } +} + +impl GpuExecutionPlan for GpuUploadExec { + fn gpu_execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + // TODO(gpu-v1): pull host batches from `self.child.execute(...)`, copy + // each column's `Buffer` into a `GpuBuffer` via `cudarc`, account the + // allocation against the task's `GpuMemoryPool`, and emit + // `GpuRecordBatch`es over a `GpuRecordBatchStream` adapter. + unimplemented!("GpuUploadExec::gpu_execute: CUDA backend not yet implemented") + } +} + +/// GPU → CPU boundary. Pulls device batches from a [`GpuExecutionPlan`] child +/// and copies them back into host [`RecordBatch`]es. +/// +/// Stores the child twice: once typed as `Arc` (so +/// `gpu_execute` can be called) and once upcast to `Arc` +/// (so the framework can borrow it via [`ExecutionPlan::children`]). Both +/// `Arc`s point to the same allocation. +pub struct GpuDownloadExec { + gpu_child: Arc, + child_as_exec: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl GpuDownloadExec { + pub fn new(gpu_child: Arc) -> Self { + let gpu_child_clone: Arc = Arc::clone(&gpu_child); + // Trait-upcast `Arc` → `Arc` + // (stable since Rust 1.86); same allocation, different vtable view. + let child_as_exec: Arc = gpu_child_clone; + let schema = child_as_exec.schema(); + let properties = passthrough_plan_properties(Arc::clone(&schema)); + Self { + gpu_child, + child_as_exec, + schema, + properties, + } + } + + pub fn gpu_input(&self) -> &Arc { + &self.gpu_child + } +} + +impl fmt::Debug for GpuDownloadExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("GpuDownloadExec") + .field("schema", &self.schema) + .finish() + } +} + +impl DisplayAs for GpuDownloadExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "GpuDownloadExec") + } +} + +impl ExecutionPlan for GpuDownloadExec { + fn name(&self) -> &str { + "GpuDownloadExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child_as_exec] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("GpuDownloadExec must have exactly one child"); + } + let [new_child] = children.try_into().unwrap(); + if Arc::ptr_eq(&new_child, &self.child_as_exec) { + return Ok(self); + } + // The framework gave us a rewritten child. We need it typed as + // `Arc` to plug into a new `GpuDownloadExec`. + // Recover that by downcasting against the closed set of GPU operator + // types and re-typing the new child via its constructor pattern. For + // v1 we only support the no-op case (pointer-equal child) — any + // other replacement would require materialising a typed Arc, which + // in turn requires each GPU op to be `Clone`. Defer. + let _typed: &dyn GpuExecutionPlan = as_gpu_execution_plan(&new_child) + .ok_or_else(|| { + datafusion_common::DataFusionError::Internal( + "GpuDownloadExec: new child is not a known GpuExecutionPlan".into(), + ) + })?; + internal_err!( + "GpuDownloadExec::with_new_children: rewriting the GPU child is not yet \ + supported (see TODO in datafusion-gpu/src/boundary.rs)" + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + // TODO(gpu-v1): call `self.gpu_child.gpu_execute(...)`, copy each + // emitted `GpuRecordBatch` from device memory back into host buffers, + // and yield `RecordBatch`es via a `RecordBatchStreamAdapter`. + unimplemented!("GpuDownloadExec::execute: CUDA backend not yet implemented") + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } +} + +// Force `Any` to be in scope so `as_any()` calls in this module resolve +// against the blanket impl when needed. +const _: fn() = || { + fn assert_any() {} + assert_any::(); + assert_any::(); +}; diff --git a/datafusion/gpu/src/execution_plan.rs b/datafusion/gpu/src/execution_plan.rs new file mode 100644 index 0000000000000..086fab0018cde --- /dev/null +++ b/datafusion/gpu/src/execution_plan.rs @@ -0,0 +1,77 @@ +// 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. + +//! [`GpuExecutionPlan`] — subtrait of `ExecutionPlan` that produces device +//! batches from `gpu_execute`. + +use std::any::Any; +use std::sync::Arc; + +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_physical_plan::ExecutionPlan; + +use crate::boundary::GpuUploadExec; +use crate::operators::{GpuFilterExec, GpuProjectionExec}; +use crate::stream::SendableGpuRecordBatchStream; + +/// An [`ExecutionPlan`] whose output is a stream of device-resident batches. +/// +/// Every `GpuExecutionPlan` is also an `ExecutionPlan` (subtrait), which gives +/// callers free trait-upcasting and lets GPU operators sit in the regular +/// physical-plan tree. The contract is that `ExecutionPlan::execute` on a +/// `GpuExecutionPlan` is unreachable in well-formed plans: the +/// [`crate::optimizer::InsertGpuExec`] rule guarantees every `GpuExecutionPlan` +/// subtree is terminated by a [`crate::boundary::GpuDownloadExec`] that calls +/// `gpu_execute` directly. +pub trait GpuExecutionPlan: ExecutionPlan { + fn gpu_execute( + &self, + partition: usize, + context: Arc, + ) -> Result; +} + +/// Borrow an `Arc` as `&dyn GpuExecutionPlan` if its +/// concrete type is one of the GPU operators defined in this crate. +/// +/// Rust trait objects don't carry subtrait vtables, so there is no built-in +/// way to recover `dyn GpuExecutionPlan`-ness from an `Arc` +/// even though every `GpuExecutionPlan` *is* an `ExecutionPlan`. Because we +/// own every concrete `GpuExecutionPlan` implementation in this crate, an +/// exhaustive `Any`-based downcast is sound. +/// +/// Callers needing an owned typed handle should hold onto the original +/// `Arc` from construction; this helper only yields a +/// borrowed view. +pub fn as_gpu_execution_plan( + plan: &Arc, +) -> Option<&dyn GpuExecutionPlan> { + // `ExecutionPlan: Any`, so we can upcast `&dyn ExecutionPlan` to `&dyn Any` + // via stable trait-upcasting coercion (Rust 1.86+). + let any: &dyn Any = plan.as_ref(); + if let Some(p) = any.downcast_ref::() { + return Some(p); + } + if let Some(p) = any.downcast_ref::() { + return Some(p); + } + if let Some(p) = any.downcast_ref::() { + return Some(p); + } + None +} diff --git a/datafusion/gpu/src/lib.rs b/datafusion/gpu/src/lib.rs new file mode 100644 index 0000000000000..449005798cbbc --- /dev/null +++ b/datafusion/gpu/src/lib.rs @@ -0,0 +1,60 @@ +// 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. + +//! Experimental GPU physical operators for DataFusion. +//! +//! This crate scaffolds a path for DataFusion logical plans to compile down +//! to physical operators that execute on Nvidia GPUs. Data flowing between +//! GPU operators lives in device memory as [`GpuRecordBatch`]es; the only +//! places where host and device data meet are the [`GpuUploadExec`] and +//! [`GpuDownloadExec`] boundary operators. +//! +//! The crate is structured around three contracts: +//! +//! * [`GpuRecordBatch`] / [`GpuRecordBatchStream`] — the device-resident +//! equivalents of `RecordBatch` and `RecordBatchStream`. +//! * [`GpuExecutionPlan`] — subtrait of `ExecutionPlan` whose `gpu_execute` +//! produces a stream of `GpuRecordBatch`es. +//! * [`GpuMemoryPool`] — device-memory accounting analogous to the host +//! `MemoryPool` in `datafusion-execution`. +//! +//! The [`InsertGpuExec`] physical optimizer rule rewrites CPU subtrees into +//! GPU ones and inserts the boundary operators automatically. +//! +//! [`GpuUploadExec`]: boundary::GpuUploadExec +//! [`GpuDownloadExec`]: boundary::GpuDownloadExec +//! [`InsertGpuExec`]: optimizer::InsertGpuExec + +#![cfg_attr(docsrs, feature(doc_cfg))] +#![deny(clippy::clone_on_ref_ptr)] +#![cfg_attr(test, allow(clippy::needless_pass_by_value))] + +pub mod batch; +pub mod boundary; +pub mod execution_plan; +pub mod memory_pool; +pub mod operators; +pub mod optimizer; +pub mod stream; + +pub use batch::{GpuBuffer, GpuRecordBatch}; +pub use execution_plan::{GpuExecutionPlan, as_gpu_execution_plan}; +pub use memory_pool::{ + GpuMemoryConsumer, GpuMemoryPool, GpuMemoryReservation, UnboundedGpuMemoryPool, +}; +pub use optimizer::InsertGpuExec; +pub use stream::{GpuRecordBatchStream, SendableGpuRecordBatchStream}; diff --git a/datafusion/gpu/src/memory_pool.rs b/datafusion/gpu/src/memory_pool.rs new file mode 100644 index 0000000000000..ffb1f9f14206d --- /dev/null +++ b/datafusion/gpu/src/memory_pool.rs @@ -0,0 +1,231 @@ +// 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. + +//! Device-memory accounting. Mirrors the shape of +//! `datafusion_execution::memory_pool::{MemoryPool, MemoryConsumer, +//! MemoryReservation}` but tracks GPU memory separately from host memory. + +use std::fmt::Debug; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion_common::{DataFusionError, Result}; + +/// Device-memory pool. Implementations decide whether and how to enforce a +/// limit. +pub trait GpuMemoryPool: Send + Sync + Debug { + fn name(&self) -> &str; + + /// Attempt to grow `reservation` by `additional` bytes. On failure the + /// reservation must be unchanged. + fn try_grow( + &self, + reservation: &GpuMemoryReservation, + additional: usize, + ) -> Result<()>; + + /// Release `bytes` from `reservation`. Must be infallible. + fn shrink(&self, reservation: &GpuMemoryReservation, bytes: usize); + + /// Total bytes currently reserved across all consumers. + fn reserved(&self) -> usize; +} + +/// Named handle used to identify an allocator for diagnostics +/// (e.g. `"GpuHashJoin:partition-3:build"`). +#[derive(Debug)] +pub struct GpuMemoryConsumer { + name: String, +} + +impl GpuMemoryConsumer { + pub fn new(name: impl Into) -> Self { + Self { name: name.into() } + } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn register(self, pool: Arc) -> GpuMemoryReservation { + GpuMemoryReservation { + consumer: self, + pool, + size: AtomicUsize::new(0), + } + } +} + +/// A unit of device-memory reservation. Dropping the reservation releases its +/// remaining bytes back to the pool. +#[derive(Debug)] +pub struct GpuMemoryReservation { + consumer: GpuMemoryConsumer, + pool: Arc, + size: AtomicUsize, +} + +impl GpuMemoryReservation { + pub fn try_grow(&self, additional: usize) -> Result<()> { + self.pool.try_grow(self, additional)?; + self.size.fetch_add(additional, Ordering::SeqCst); + Ok(()) + } + + pub fn shrink(&self, bytes: usize) { + let cur = self.size.load(Ordering::SeqCst); + let to_release = bytes.min(cur); + self.pool.shrink(self, to_release); + self.size.fetch_sub(to_release, Ordering::SeqCst); + } + + pub fn size(&self) -> usize { + self.size.load(Ordering::SeqCst) + } + + pub fn consumer(&self) -> &GpuMemoryConsumer { + &self.consumer + } +} + +impl Drop for GpuMemoryReservation { + fn drop(&mut self) { + let remaining = self.size.load(Ordering::SeqCst); + if remaining > 0 { + self.pool.shrink(self, remaining); + } + } +} + +/// Pool with no upper limit; tracks usage for diagnostics only. Useful for +/// tests and for the v1 scaffold before a CUDA-backed pool exists. +#[derive(Debug, Default)] +pub struct UnboundedGpuMemoryPool { + reserved: AtomicUsize, +} + +impl UnboundedGpuMemoryPool { + pub fn new() -> Self { + Self::default() + } +} + +impl GpuMemoryPool for UnboundedGpuMemoryPool { + fn name(&self) -> &str { + "UnboundedGpuMemoryPool" + } + + fn try_grow( + &self, + _reservation: &GpuMemoryReservation, + additional: usize, + ) -> Result<()> { + self.reserved.fetch_add(additional, Ordering::SeqCst); + Ok(()) + } + + fn shrink(&self, _reservation: &GpuMemoryReservation, bytes: usize) { + self.reserved.fetch_sub(bytes, Ordering::SeqCst); + } + + fn reserved(&self) -> usize { + self.reserved.load(Ordering::SeqCst) + } +} + +/// Bounded pool with a hard byte cap; `try_grow` returns +/// [`DataFusionError::ResourcesExhausted`] when the cap would be exceeded. +#[derive(Debug)] +pub struct BoundedGpuMemoryPool { + limit: usize, + reserved: AtomicUsize, +} + +impl BoundedGpuMemoryPool { + pub fn new(limit: usize) -> Self { + Self { + limit, + reserved: AtomicUsize::new(0), + } + } + + pub fn limit(&self) -> usize { + self.limit + } +} + +impl GpuMemoryPool for BoundedGpuMemoryPool { + fn name(&self) -> &str { + "BoundedGpuMemoryPool" + } + + fn try_grow( + &self, + reservation: &GpuMemoryReservation, + additional: usize, + ) -> Result<()> { + let prev = self.reserved.fetch_add(additional, Ordering::SeqCst); + if prev + additional > self.limit { + self.reserved.fetch_sub(additional, Ordering::SeqCst); + return Err(DataFusionError::ResourcesExhausted(format!( + "GPU memory limit exceeded: consumer {:?} requested {} bytes \ + but pool {} has only {} bytes free (limit {})", + reservation.consumer().name(), + additional, + self.name(), + self.limit.saturating_sub(prev), + self.limit, + ))); + } + Ok(()) + } + + fn shrink(&self, _reservation: &GpuMemoryReservation, bytes: usize) { + self.reserved.fetch_sub(bytes, Ordering::SeqCst); + } + + fn reserved(&self) -> usize { + self.reserved.load(Ordering::SeqCst) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reservation_releases_on_drop() { + let pool: Arc = Arc::new(UnboundedGpuMemoryPool::new()); + { + let reservation = GpuMemoryConsumer::new("test").register(Arc::clone(&pool)); + reservation.try_grow(1024).unwrap(); + assert_eq!(pool.reserved(), 1024); + } + assert_eq!(pool.reserved(), 0); + } + + #[test] + fn bounded_pool_rejects_over_limit() { + let pool: Arc = Arc::new(BoundedGpuMemoryPool::new(1024)); + let reservation = GpuMemoryConsumer::new("test").register(Arc::clone(&pool)); + reservation.try_grow(512).unwrap(); + let err = reservation.try_grow(1024).unwrap_err(); + assert!(matches!(err, DataFusionError::ResourcesExhausted(_))); + assert_eq!(reservation.size(), 512); + assert_eq!(pool.reserved(), 512); + } +} diff --git a/datafusion/gpu/src/operators.rs b/datafusion/gpu/src/operators.rs new file mode 100644 index 0000000000000..8ff6ea69e0fa8 --- /dev/null +++ b/datafusion/gpu/src/operators.rs @@ -0,0 +1,290 @@ +// 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. + +//! v1 GPU operator stubs: [`GpuProjectionExec`] and [`GpuFilterExec`]. +//! +//! Both store a `GpuExecutionPlan` child (typed) and an upcast +//! `Arc` view of the same allocation so the framework can +//! borrow children via `ExecutionPlan::children`. Bodies are stubbed pending +//! the CUDA kernel work. + +use std::fmt; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, PlanProperties, +}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, +}; + +use crate::execution_plan::GpuExecutionPlan; +use crate::stream::SendableGpuRecordBatchStream; + +fn projection_plan_properties(schema: SchemaRef) -> Arc { + Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )) +} + +/// GPU equivalent of `ProjectionExec`. v1 plan: column passthrough where the +/// expression is `Column`; arithmetic via a hand-rolled CUDA kernel. +pub struct GpuProjectionExec { + expressions: Vec<(Arc, String)>, + gpu_input: Arc, + input_as_exec: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl GpuProjectionExec { + pub fn try_new( + expressions: Vec<(Arc, String)>, + gpu_input: Arc, + schema: SchemaRef, + ) -> Result { + let gpu_input_clone: Arc = Arc::clone(&gpu_input); + let input_as_exec: Arc = gpu_input_clone; + let properties = projection_plan_properties(Arc::clone(&schema)); + Ok(Self { + expressions, + gpu_input, + input_as_exec, + schema, + properties, + }) + } + + pub fn expressions(&self) -> &[(Arc, String)] { + &self.expressions + } + + pub fn gpu_input(&self) -> &Arc { + &self.gpu_input + } +} + +impl fmt::Debug for GpuProjectionExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("GpuProjectionExec") + .field("schema", &self.schema) + .finish() + } +} + +impl DisplayAs for GpuProjectionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "GpuProjectionExec") + } +} + +impl ExecutionPlan for GpuProjectionExec { + fn name(&self) -> &str { + "GpuProjectionExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input_as_exec] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("GpuProjectionExec must have exactly one child"); + } + let [new_child] = children.try_into().unwrap(); + if Arc::ptr_eq(&new_child, &self.input_as_exec) { + return Ok(self); + } + internal_err!( + "GpuProjectionExec::with_new_children: replacing the GPU input is not \ + yet supported in v1 scaffold" + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + internal_err!( + "GpuProjectionExec produces device-resident batches; wrap in \ + GpuDownloadExec before consuming on the host" + ) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for (expr, _alias) in &self.expressions { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + Ok(tnr) + } +} + +impl GpuExecutionPlan for GpuProjectionExec { + fn gpu_execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + // TODO(gpu-v1): walk `self.expressions`; for each `Column`, return the + // input column unchanged; for arithmetic, dispatch to the appropriate + // CUDA kernel; assemble a new `GpuRecordBatch` against the projected + // schema. + unimplemented!("GpuProjectionExec::gpu_execute: CUDA kernels not yet implemented") + } +} + +/// GPU equivalent of `FilterExec`. v1 plan: predicate is evaluated as a +/// boolean mask kernel, then a gather kernel materialises the surviving rows. +pub struct GpuFilterExec { + predicate: Arc, + gpu_input: Arc, + input_as_exec: Arc, + schema: SchemaRef, + properties: Arc, +} + +impl GpuFilterExec { + pub fn try_new( + predicate: Arc, + gpu_input: Arc, + ) -> Result { + let gpu_input_clone: Arc = Arc::clone(&gpu_input); + let input_as_exec: Arc = gpu_input_clone; + let schema = input_as_exec.schema(); + let properties = projection_plan_properties(Arc::clone(&schema)); + Ok(Self { + predicate, + gpu_input, + input_as_exec, + schema, + properties, + }) + } + + pub fn predicate(&self) -> &Arc { + &self.predicate + } + + pub fn gpu_input(&self) -> &Arc { + &self.gpu_input + } +} + +impl fmt::Debug for GpuFilterExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("GpuFilterExec") + .field("schema", &self.schema) + .finish() + } +} + +impl DisplayAs for GpuFilterExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "GpuFilterExec") + } +} + +impl ExecutionPlan for GpuFilterExec { + fn name(&self) -> &str { + "GpuFilterExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input_as_exec] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!("GpuFilterExec must have exactly one child"); + } + let [new_child] = children.try_into().unwrap(); + if Arc::ptr_eq(&new_child, &self.input_as_exec) { + return Ok(self); + } + internal_err!( + "GpuFilterExec::with_new_children: replacing the GPU input is not yet \ + supported in v1 scaffold" + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + internal_err!( + "GpuFilterExec produces device-resident batches; wrap in \ + GpuDownloadExec before consuming on the host" + ) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + f(self.predicate.as_ref())?; + Ok(TreeNodeRecursion::Continue) + } +} + +impl GpuExecutionPlan for GpuFilterExec { + fn gpu_execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + // TODO(gpu-v1): evaluate `self.predicate` against the input batch on + // device, producing a boolean selection mask; gather surviving rows + // into a new `GpuRecordBatch`. + unimplemented!("GpuFilterExec::gpu_execute: CUDA kernels not yet implemented") + } +} diff --git a/datafusion/gpu/src/optimizer.rs b/datafusion/gpu/src/optimizer.rs new file mode 100644 index 0000000000000..7366cf4e3c300 --- /dev/null +++ b/datafusion/gpu/src/optimizer.rs @@ -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. + +//! [`InsertGpuExec`] — physical optimizer rule that rewrites supported CPU +//! subtrees into GPU operators with [`crate::boundary::GpuUploadExec`] / +//! [`crate::boundary::GpuDownloadExec`] at the boundaries. + +use std::sync::Arc; + +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_physical_optimizer::optimizer::PhysicalOptimizerRule; +use datafusion_physical_plan::ExecutionPlan; + +/// Inserts GPU operators into a physical plan. +/// +/// In the v1 scaffold this rule is a no-op: it returns the plan unchanged. +/// The pattern-matching logic that swaps `ProjectionExec` / `FilterExec` +/// subtrees for their GPU equivalents lands once the CUDA kernels are wired +/// up. The rule is shipped now so it can be registered against a +/// `SessionStateBuilder` end-to-end and so its position in the rule order can +/// be validated against the rest of the physical optimizer pipeline. +#[derive(Debug, Default)] +pub struct InsertGpuExec; + +impl InsertGpuExec { + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for InsertGpuExec { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + // TODO(gpu-v1): walk the plan tree; for maximal subtrees consisting + // only of `ProjectionExec` and `FilterExec` over a non-GPU leaf, + // produce a GPU subtree of `GpuProjectionExec`/`GpuFilterExec` rooted + // under a `GpuDownloadExec` and bottomed by a `GpuUploadExec` wrapping + // the original leaf. + Ok(plan) + } + + fn name(&self) -> &str { + "InsertGpuExec" + } + + fn schema_check(&self) -> bool { + true + } +} diff --git a/datafusion/gpu/src/stream.rs b/datafusion/gpu/src/stream.rs new file mode 100644 index 0000000000000..f23624218ab86 --- /dev/null +++ b/datafusion/gpu/src/stream.rs @@ -0,0 +1,37 @@ +// 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. + +//! Async stream of [`GpuRecordBatch`]es. Parallel to `RecordBatchStream` in +//! `datafusion-execution`, intentionally a distinct type so device-resident +//! batches cannot accidentally flow into CPU operators. + +use std::pin::Pin; + +use arrow::datatypes::SchemaRef; +use datafusion_common::Result; +use futures::Stream; + +use crate::batch::GpuRecordBatch; + +/// Async stream of device-resident batches with a known schema. +pub trait GpuRecordBatchStream: Stream> { + fn schema(&self) -> SchemaRef; +} + +/// A pinned, sendable [`GpuRecordBatchStream`]. The return type of +/// [`crate::GpuExecutionPlan::gpu_execute`]. +pub type SendableGpuRecordBatchStream = Pin>;