Added aggregate expr and kernel

Author: manishkr

.github/actions/setup/action.yml

@@ -31,9 +31,10 @@ runs:
 
         - name: Setup CUDA
           run: |
-              wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
-              dpkg -i cuda-keyring_1.1-1_all.deb
-              apt-get update
-              apt-get install -y cuda-toolkit-13-0
-              apt-get install -y cuda-drivers
+              # Fix when cuda keyring certifcate is fixed
+              #wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+              #dpkg -i cuda-keyring_1.1-1_all.deb
+              #apt-get update
+              #apt-get install -y cuda-toolkit-13-0
+              #apt-get install -y cuda-drivers
           shell: bash

tachyon/compute/Cargo.toml

@@ -31,3 +31,7 @@ criterion.workspace = true
 [[bench]]
 name = "evaluate_bench"
 harness = false
+
+[[bench]]
+name = "aggregate_bench"
+harness = false

tachyon/compute/benches/aggregate_bench.rs

@@ -0,0 +1,121 @@
+use std::hint::black_box;
+use std::mem::size_of;
+
+use compute::bit_vector::BitVector;
+use compute::column::{Column, VecArray};
+use compute::data_type::DataType;
+use compute::error::ErrorMode;
+use compute::evaluate::{Device, evaluate};
+use compute::expr::Expr;
+use compute::operator::Operator;
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use rand::Rng;
+use tokio::runtime::Builder;
+
+fn single_thread_runtime() -> tokio::runtime::Runtime {
+    Builder::new_current_thread().enable_all().build().unwrap()
+}
+
+fn random_i32(len: usize) -> Vec<i32> {
+    let mut rng = rand::rng();
+    (0..len).map(|_| rng.random_range(i32::MIN..i32::MAX)).collect::<Vec<i32>>()
+}
+
+fn random_f64(len: usize) -> Vec<f64> {
+    let mut rng = rand::rng();
+    (0..len).map(|_| rng.random_range(-1_000_000.0..1_000_000.0)).collect::<Vec<f64>>()
+}
+
+fn make_i32_column(name: &str, len: usize) -> Column<u64> {
+    use std::sync::Arc;
+    let values = Arc::new(VecArray { data: random_i32(len), datatype: DataType::I32 });
+    Column::new(name, values, Some(BitVector::<u64>::new_all_valid(len)))
+}
+
+fn make_f64_column(name: &str, len: usize) -> Column<u64> {
+    use std::sync::Arc;
+    let values = Arc::new(VecArray { data: random_f64(len), datatype: DataType::F64 });
+    Column::new(name, values, Some(BitVector::<u64>::new_all_valid(len)))
+}
+
+fn bench_aggregate_i32(c: &mut Criterion, rt: &tokio::runtime::Runtime, len: usize) {
+    let col_a = make_i32_column("a", len);
+    let ops = [Operator::Min, Operator::Max, Operator::Sum, Operator::Count];
+    let bytes = (len * size_of::<i32>()) as u64;
+
+    let mut group_rows = c.benchmark_group("aggregate_i32_rows");
+    group_rows.throughput(Throughput::Elements(len as u64));
+    for op in ops {
+        let expr = Expr::aggregate(op, Expr::col("a"), false);
+        let bench_id = BenchmarkId::new(format!("{:?}", op), len);
+        group_rows.bench_with_input(bench_id, &op, |bch, &_op| {
+            bch.iter(|| {
+                let out = rt.block_on(async {
+                    evaluate(Device::GPU, ErrorMode::Tachyon, &expr, black_box(&[col_a.clone()]))
+                        .await
+                });
+                black_box(out).unwrap();
+            });
+        });
+    }
+    group_rows.finish();
+
+    let mut group_bytes = c.benchmark_group("aggregate_i32_bytes");
+    group_bytes.throughput(Throughput::Bytes(bytes));
+    for op in ops {
+        let expr = Expr::aggregate(op, Expr::col("a"), false);
+        let bench_id = BenchmarkId::new(format!("{:?}", op), len);
+        group_bytes.bench_with_input(bench_id, &op, |bch, &_op| {
+            bch.iter(|| {
+                let out = rt.block_on(async {
+                    evaluate(Device::GPU, ErrorMode::Tachyon, &expr, black_box(&[col_a.clone()]))
+                        .await
+                });
+                black_box(out).unwrap();
+            });
+        });
+    }
+    group_bytes.finish();
+}
+
+fn bench_aggregate_f64(c: &mut Criterion, rt: &tokio::runtime::Runtime, len: usize) {
+    let col_a = make_f64_column("a", len);
+    let expr = Expr::aggregate(Operator::Avg, Expr::col("a"), false);
+    let bytes = (len * size_of::<f64>()) as u64;
+
+    let mut group_rows = c.benchmark_group("aggregate_f64_rows");
+    group_rows.throughput(Throughput::Elements(len as u64));
+    let bench_id_rows = BenchmarkId::new("Avg", len);
+    group_rows.bench_with_input(bench_id_rows, &len, |bch, &_len| {
+        bch.iter(|| {
+            let out = rt.block_on(async {
+                evaluate(Device::GPU, ErrorMode::Tachyon, &expr, black_box(&[col_a.clone()])).await
+            });
+            black_box(out).unwrap();
+        });
+    });
+    group_rows.finish();
+
+    let mut group_bytes = c.benchmark_group("aggregate_f64_bytes");
+    group_bytes.throughput(Throughput::Bytes(bytes));
+    let bench_id_bytes = BenchmarkId::new("Avg", len);
+    group_bytes.bench_with_input(bench_id_bytes, &len, |bch, &_len| {
+        bch.iter(|| {
+            let out = rt.block_on(async {
+                evaluate(Device::GPU, ErrorMode::Tachyon, &expr, black_box(&[col_a.clone()])).await
+            });
+            black_box(out).unwrap();
+        });
+    });
+    group_bytes.finish();
+}
+
+fn bench_all_aggregate(c: &mut Criterion) {
+    let rt = single_thread_runtime();
+    let len = 1_000_000;
+    bench_aggregate_i32(c, &rt, len);
+    bench_aggregate_f64(c, &rt, len);
+}
+
+criterion_group!(benches, bench_all_aggregate);
+criterion_main!(benches);

tachyon/compute/src/codegen.rs

@@ -232,6 +232,12 @@ impl CodeGen for Expr {
                     });
                 var
             }
+            Expr::Aggregate { op, .. } => {
+                return Err(TypeError::Unsupported(format!(
+                    "Aggregate {:?} is not supported in codegen yet",
+                    op
+                )));
+            }
         };
 
         Ok(var)

tachyon/compute/src/evaluate.rs

@@ -16,6 +16,7 @@ use crate::column::Column;
 use crate::data_type::DataType;
 use crate::error::ErrorMode;
 use crate::expr::{Expr, SchemaContext};
+use crate::operator::Operator;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Device {
@@ -40,15 +41,27 @@ async fn evaluate_gpu<B: BitBlock>(
         .collect();
 
     let schema_context = SchemaContext::new().with_columns(&column_map).with_error_mode(error_mode);
+    match expr {
+        Expr::Aggregate { op, arg, distinct } => {
+            evaluate_gpu_aggregate::<B>(*op, arg.as_ref(), *distinct, &schema_context, columns)
+                .await
+        }
+        _ => evaluate_gpu_row::<B>(expr, &schema_context, columns).await,
+    }
+}
+
+async fn evaluate_gpu_row<B: BitBlock>(
+    expr: &Expr, schema_context: &SchemaContext, columns: &[Column<B>],
+) -> Result<Vec<Column<B>>, Box<dyn Error>> {
     let mut code_block = CodeBlock::default();
-    expr.to_nvrtc::<B>(&schema_context, &mut code_block)?;
+    expr.to_nvrtc::<B>(schema_context, &mut code_block)?;
 
     let size = columns[0].len();
     let input_cols =
         columns.iter().map(|col| col.to_gpu_column()).collect::<Result<Vec<_>, _>>()?;
 
     let mut output_cols = Vec::<gpu_column::Column>::new();
-    let result_type = expr.infer_type(&schema_context)?;
+    let result_type = expr.infer_type(schema_context)?;
 
     let gpu_col = gpu_column::Column::new_uninitialized::<B>(
         size * result_type.native_size(),
@@ -68,3 +81,80 @@ async fn evaluate_gpu<B: BitBlock>(
 
     Ok(result_cols)
 }
+
+async fn evaluate_gpu_aggregate<B: BitBlock>(
+    op: Operator, arg: &Expr, distinct: bool, schema_context: &SchemaContext, columns: &[Column<B>],
+) -> Result<Vec<Column<B>>, Box<dyn Error>> {
+    if distinct {
+        return Err("DISTINCT aggregates are not supported yet".into());
+    }
+
+    let (col_idx, col_type) = match arg {
+        Expr::Column(col_name) => schema_context
+            .lookup(col_name)
+            .copied()
+            .ok_or_else(|| format!("unknown column: {}", col_name))?,
+        _ => return Err("Aggregate argument must be a column reference".into()),
+    };
+
+    let result_type =
+        Expr::Aggregate { op, arg: Box::new(arg.clone()), distinct }.infer_type(schema_context)?;
+    let code = build_aggregate_nvrtc_code::<B>(
+        op,
+        col_idx,
+        col_type,
+        result_type,
+        schema_context.error_mode() == ErrorMode::Ansi,
+    )?;
+    let input_cols =
+        columns.iter().map(|col| col.to_gpu_column()).collect::<Result<Vec<_>, _>>()?;
+
+    let mut output_cols = Vec::<gpu_column::Column>::new();
+    let size = 1usize;
+    let gpu_col = gpu_column::Column::new_uninitialized::<B>(
+        size * result_type.native_size(),
+        size.div_ceil(B::BITS),
+        size,
+    )?;
+    output_cols.push(gpu_col);
+
+    cuda_launcher::launch_aggregate::<B>(&code, &input_cols, &output_cols).await?;
+
+    let result_cols = output_cols
+        .into_iter()
+        .map(|col| -> Result<_, Box<dyn Error>> {
+            Column::from_gpu_column(&col, "r0", result_type)
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    Ok(result_cols)
+}
+
+fn build_aggregate_nvrtc_code<B: BitBlock>(
+    op: Operator, col_idx: u16, col_type: DataType, result_type: DataType, ansi_error_mode: bool,
+) -> Result<String, Box<dyn Error>> {
+    let input_kernel_type = col_type.kernel_type();
+    let output_kernel_type = result_type.kernel_type();
+    let bits_type = B::C_TYPE;
+
+    let code = match op {
+        Operator::Min => format!(
+            "\taggregate_codegen::min<TypeKind::{input_kernel_type}, TypeKind::{output_kernel_type}, {bits_type}>(ctx, input, output, num_rows, {col_idx});\n"
+        ),
+        Operator::Max => format!(
+            "\taggregate_codegen::max<TypeKind::{input_kernel_type}, TypeKind::{output_kernel_type}, {bits_type}>(ctx, input, output, num_rows, {col_idx});\n"
+        ),
+        Operator::Sum => format!(
+            "\taggregate_codegen::sum<{ansi_error_mode}, TypeKind::{input_kernel_type}, TypeKind::{output_kernel_type}, {bits_type}>(ctx, input, output, num_rows, {col_idx});\n"
+        ),
+        Operator::Avg => format!(
+            "\taggregate_codegen::avg<{ansi_error_mode}, TypeKind::{input_kernel_type}, TypeKind::{output_kernel_type}, {bits_type}>(ctx, input, output, num_rows, {col_idx});\n"
+        ),
+        Operator::Count => format!(
+            "\taggregate_codegen::count<TypeKind::{input_kernel_type}, {bits_type}>(ctx, input, output, num_rows, {col_idx});\n"
+        ),
+        _ => return Err(format!("Unsupported aggregate operator: {:?}", op).into()),
+    };
+
+    Ok(code)
+}

tachyon/compute/src/expr.rs

@@ -12,7 +12,7 @@ use half::{bf16, f16};
 
 use crate::data_type::DataType;
 use crate::error::ErrorMode;
-use crate::operator::Operator;
+pub use crate::operator::Operator;
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum Literal {
@@ -47,6 +47,8 @@ pub enum Expr {
     Call { name: String, args: Vec<Expr> },
 
     Cast { expr: Box<Expr>, to: DataType },
+
+    Aggregate { op: Operator, arg: Box<Expr>, distinct: bool },
 }
 
 impl Expr {
@@ -126,6 +128,10 @@ impl Expr {
         Expr::Cast { expr: Box::new(self), to }
     }
 
+    pub fn aggregate(op: Operator, arg: Expr, distinct: bool) -> Self {
+        Expr::Aggregate { op, arg: Box::new(arg), distinct }
+    }
+
     pub fn children(&self) -> Vec<&Expr> {
         match self {
             Expr::Column(_) | Expr::Literal(_) => vec![],
@@ -134,6 +140,7 @@ impl Expr {
             Expr::Nary { args, .. } => args.iter().map(|x| x.as_ref()).collect(),
             Expr::Call { args, .. } => args.iter().collect(),
             Expr::Cast { expr, .. } => vec![expr.as_ref()],
+            Expr::Aggregate { arg, .. } => vec![arg.as_ref()],
         }
     }
 }
@@ -313,6 +320,31 @@ impl Expr {
                 let _ = expr.infer_type(schema)?;
                 Ok(*to)
             }
+
+            Expr::Aggregate { op, arg, .. } => {
+                let t = arg.infer_type(schema)?;
+                match op {
+                    Operator::Count => Ok(DataType::U64),
+                    Operator::Sum | Operator::Avg | Operator::Min | Operator::Max => match t {
+                        DataType::I8
+                        | DataType::I16
+                        | DataType::I32
+                        | DataType::I64
+                        | DataType::U8
+                        | DataType::U16
+                        | DataType::U32
+                        | DataType::U64
+                        | DataType::BF16
+                        | DataType::F16
+                        | DataType::F32
+                        | DataType::F64 => Ok(t),
+                        _ => Err(TypeError::Unsupported(format!("aggregate {:?} on {:?}", op, t))),
+                    },
+                    _ => {
+                        Err(TypeError::Unsupported(format!("not an aggregate operator: {:?}", op)))
+                    }
+                }
+            }
         }
     }
 
@@ -413,7 +445,17 @@ impl Expr {
                 Ok(Expr::Call { name: name.clone(), args: simplified_args })
             }
 
-            _ => Ok(self.clone()),
+            Expr::Cast { expr, to } => {
+                let simplified_expr = expr.simplify(schema)?;
+                Ok(Expr::Cast { expr: Box::new(simplified_expr), to: *to })
+            }
+
+            Expr::Aggregate { op, arg, distinct } => {
+                let simplified_arg = arg.simplify(schema)?;
+                Ok(Expr::Aggregate { op: *op, arg: Box::new(simplified_arg), distinct: *distinct })
+            }
+
+            Expr::Column(_) | Expr::Literal(_) => Ok(self.clone()),
         }
     }
 }
@@ -495,6 +537,13 @@ impl fmt::Display for Expr {
             Expr::Nary { op, args } => write!(f, "{}({:?})", op, args),
             Expr::Call { name, args } => write!(f, "{}({:?})", name, args),
             Expr::Cast { expr, to } => write!(f, "cast({} as {:?})", expr, to),
+            Expr::Aggregate { op, arg, distinct } => {
+                if *distinct {
+                    write!(f, "{}(DISTINCT {})", op, arg)
+                } else {
+                    write!(f, "{}({})", op, arg)
+                }
+            }
         }
     }
 }