support order sparse 64

Author: hero78119

Cargo.lock

@@ -1,16 +1,16 @@
 [workspace]
 members = [
-  "ceno_cli",
-  "ceno_emul",
-  "ceno_host",
-  "ceno_serde",
-  "ceno_rt",
-  "ceno_zkvm",
-  "ceno_recursion",
-  "derive",
-  "examples-builder",
-  "examples",
-  "guest_libs/*",
+    "ceno_cli",
+    "ceno_emul",
+    "ceno_host",
+    "ceno_serde",
+    "ceno_rt",
+    "ceno_zkvm",
+    "ceno_recursion",
+    "derive",
+    "examples-builder",
+    "examples",
+    "guest_libs/*",
 ]
 resolver = "2"
 
@@ -66,11 +66,11 @@ secp = "0.4.1"
 serde = { version = "1.0", features = ["derive", "rc"] }
 serde_json = "1.0"
 smallvec = { version = "1.13.2", features = [
-  "const_generics",
-  "const_new",
-  "serde",
-  "union",
-  "write",
+    "const_generics",
+    "const_new",
+    "serde",
+    "union",
+    "write",
 ] }
 strum = "0.26"
 strum_macros = "0.26"
@@ -79,7 +79,7 @@ thiserror = "2"
 thread_local = "1.1"
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tracing = { version = "0.1", features = [
-  "attributes",
+    "attributes",
 ] }
 tracing-forest = { version = "0.1.6" }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
@@ -127,8 +127,8 @@ lto = "thin"
 #ceno_crypto_primitives = { path = "../ceno-patch/crypto-primitives", package = "ceno_crypto_primitives" }
 #ceno_syscall = { path = "../ceno-patch/syscall", package = "ceno_syscall" }
 
-#[patch."https://github.com/scroll-tech/ceno-gpu-mock.git"]
-#ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal", default-features = false, features = ["bb31"] }
+[patch."https://github.com/scroll-tech/ceno-gpu-mock.git"]
+ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal", default-features = false, features = ["bb31"] }
 
 #[patch."https://github.com/scroll-tech/gkr-backend"]
 #ff_ext = { path = "../gkr-backend/crates/ff_ext", package = "ff_ext" }

Cargo.toml

@@ -54,9 +54,9 @@ path = "src/bin/e2e_aggregate.rs"
 [features]
 bench-metrics = ["openvm-circuit/metrics"]
 default = ["parallel", "nightly-features"]
-gpu = ["openvm-circuit/cuda", "openvm-native-circuit/cuda", "dep:openvm-cuda-backend"]
+gpu = ["ceno_zkvm/gpu", "openvm-circuit/cuda", "openvm-native-circuit/cuda", "dep:openvm-cuda-backend"]
 nightly-features = [
-  "ceno_zkvm/nightly-features",
-  "openvm-stark-sdk/nightly-features",
+    "ceno_zkvm/nightly-features",
+    "openvm-stark-sdk/nightly-features",
 ]
 parallel = ["openvm-stark-backend/parallel"]

ceno_recursion/Cargo.toml

@@ -129,21 +129,18 @@ pub fn build_eq_x_r_with_sel_gpu<E: ExtensionField>(
     }
 
     let eq_len = 1 << point.len();
-    let (num_instances, is_sp32, indices) = match selector {
+    let (num_instances, sparse_num_var, indices) = match selector {
         SelectorType::None => panic!("SelectorType::None"),
-        SelectorType::Whole(_expr) => (eq_len, false, vec![]),
-        SelectorType::Prefix(_expr) => (selector_ctx.num_instances, false, vec![]),
-        SelectorType::OrderedSparse32 { indices, .. } => {
-            (selector_ctx.num_instances, true, indices.clone())
-        }
-        SelectorType::OrderedSparse64 { .. } => {
-            unimplemented!("OrderedSparse64 is not supported in GPU selector path")
-        }
+        SelectorType::Whole(_expr) => (eq_len, 0, vec![]),
+        SelectorType::Prefix(_expr) => (selector_ctx.num_instances, 0, vec![]),
+        SelectorType::OrderedSparse {
+            indices, num_vars, ..
+        } => (selector_ctx.num_instances, *num_vars, indices.clone()),
         SelectorType::QuarkBinaryTreeLessThan(..) => unimplemented!(),
     };
 
     // type eq
-    let eq_mle = if is_sp32 {
+    let eq_mle = if sparse_num_var > 0 {
         assert_eq!(selector_ctx.offset, 0);
         let eq = build_eq_x_r_gpu(hal, point);
         let mut eq_buf = match eq.mle {
@@ -152,11 +149,12 @@ pub fn build_eq_x_r_with_sel_gpu<E: ExtensionField>(
             GpuFieldType::Unreachable => panic!("Unreachable GpuFieldType"),
         };
         let indices_u32 = indices.iter().map(|x| *x as u32).collect_vec();
-        ordered_sparse32_selector_gpu::<CudaHalBB31, BB31Ext, BB31Base>(
+        ordered_sparse_selector_gpu::<CudaHalBB31, BB31Ext, BB31Base>(
             &hal.inner,
             &mut eq_buf.buf,
             &indices_u32,
             num_instances,
+            sparse_num_var,
         )
         .unwrap();
         eq_buf

gkr_iop/src/gkr/layer/gpu/utils.rs

@@ -29,7 +29,7 @@ pub mod gpu_prover {
             buffer::BufferImpl,
             get_ceno_gpu_device_id,
             mle::{
-                build_mle_as_ceno, ordered_sparse32_selector_gpu, rotation_next_base_mle_gpu,
+                build_mle_as_ceno, ordered_sparse_selector_gpu, rotation_next_base_mle_gpu,
                 rotation_selector_gpu,
             },
             utils::HasUtils,