perf: avoid cloning CPU witness MLEs (#1342)

## Problem

left-over from #923. CPU trace
commit cloned large witness MLEs before proving, adding avoidable memory
traffic on the prover hot path.



## Design Rationale

Keep committed witness MLEs behind `Arc` and drain/transport ownership
where possible, avoiding deep clones without changing proof semantics.

## Change Highlights

- `ceno_zkvm`: return `Arc` witness MLEs from trace commit and consume
structural MLEs during transport.
- `ceno_zkvm`: keep GPU trait shape aligned while preserving existing
GPU behavior.

## Benchmark / Performance Impact

### Operation

| Operation | master (s) | this PR (s) | Improve (master -> this PR) |
|-----------|------------|-------------|-----------------------------|
| CPU proving, keccak e2e shard total | 6.942 | 6.596 | 4.98% faster |
| GPU proving, keccak e2e shard total | 1.191 | 1.186 | 0.44% faster |

### Layer

| Layer | master (s) | this PR (s) | Improve (master -> this PR) |
|-------|------------|-------------|-----------------------------|
| N/A: shard-level proving total measured | N/A | N/A | no regression
observed |

Benchmark command(s):

```sh
cargo run --config net.git-fetch-with-cli=true --release --package ceno_zkvm --bin e2e -- --platform=ceno --max-cycle-per-shard=1600 examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall
cargo run --config net.git-fetch-with-cli=true --features gpu --release --package ceno_zkvm --bin e2e -- --platform=ceno --max-cycle-per-shard=1600 examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall
```

Environment: local x86_64 Linux, release build, local
`../ceno-gpu/cuda_hal` patch for GPU validation.

raw data:

- master: CPU shards `3.272s + 3.670s`; GPU shards `0.624s + 0.568s`
- this PR: CPU shards `3.336s + 3.260s`; GPU shards `0.593s + 0.593s`

## Testing

```sh
cargo check --config net.git-fetch-with-cli=true --package ceno_zkvm --bin e2e
cargo run --config net.git-fetch-with-cli=true --release --package ceno_zkvm --bin e2e -- --platform=ceno --max-cycle-per-shard=1600 examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall
cargo run --config net.git-fetch-with-cli=true --features gpu --release --package ceno_zkvm --bin e2e -- --platform=ceno --max-cycle-per-shard=1600 examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall
```

## Risks and Rollout

Low risk: prover-side ownership change only. Rollback is reverting the
`Arc` witness-MLE plumbing.

## Follow-ups (optional)

None.

## Copilot Reviewer Directive (keep this section)

When Copilot reviews this PR, apply `.github/copilot-instructions.md`
strictly.

Author: hero78119

ceno_zkvm/src/scheme/cpu/mod.rs

@@ -558,7 +558,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> TraceCommitter<CpuBa
         &self,
         traces: BTreeMap<usize, witness::RowMajorMatrix<E::BaseField>>,
     ) -> (
-        Vec<MultilinearExtension<'a, E>>,
+        Vec<ArcMultilinearExtension<'a, E>>,
         PCS::CommitmentWithWitness,
         PCS::Commitment,
     ) {
@@ -576,22 +576,19 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> TraceCommitter<CpuBa
         let prover_param = &self.backend.pp;
         let pcs_data = PCS::batch_commit(prover_param, traces.into_values().collect_vec()).unwrap();
         let commit = PCS::get_pure_commitment(&pcs_data);
-        let mles = PCS::get_arc_mle_witness_from_commitment(&pcs_data)
-            .into_par_iter()
-            .map(|mle| mle.as_ref().clone())
-            .collect::<Vec<_>>();
+        let mles = PCS::get_arc_mle_witness_from_commitment(&pcs_data);
 
         (mles, pcs_data, commit)
     }
 
     fn extract_witness_mles<'a, 'b>(
         &self,
-        witness_mles: &'b mut Vec<<CpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>,
+        witness_mles: &'b mut Vec<Arc<<CpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>>,
         _pcs_data: &'b <CpuBackend<E, PCS> as ProverBackend>::PcsData,
     ) -> Box<
         dyn Iterator<Item = Arc<<CpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>> + 'b,
     > {
-        let iter = witness_mles.drain(..).map(Arc::new);
+        let iter = witness_mles.drain(..);
         Box::new(iter)
     }
 }
@@ -1524,9 +1521,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> DeviceTransporter<Cp
 
     fn transport_mles<'a>(
         &self,
-        mles: &[MultilinearExtension<'a, E>],
+        mles: Vec<MultilinearExtension<'a, E>>,
     ) -> Vec<ArcMultilinearExtension<'a, E>> {
-        mles.iter().map(|mle| mle.clone().into()).collect_vec()
+        mles.into_iter().map(Arc::new).collect_vec()
     }
 }
 

ceno_zkvm/src/scheme/gpu/mod.rs

@@ -1670,7 +1670,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + 'static>
         &self,
         traces: BTreeMap<usize, witness::RowMajorMatrix<E::BaseField>>,
     ) -> (
-        Vec<<GpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>,
+        Vec<Arc<<GpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>>,
         <GpuBackend<E, PCS> as ProverBackend>::PcsData,
         PCS::Commitment,
     ) {
@@ -1811,7 +1811,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + 'static>
 
     fn extract_witness_mles<'a, 'b>(
         &self,
-        _witness_mles: &'b mut Vec<<GpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>,
+        _witness_mles: &'b mut Vec<Arc<<GpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>>,
         pcs_data: &'b <GpuBackend<E, PCS> as ProverBackend>::PcsData,
     ) -> Box<
         dyn Iterator<Item = Arc<<GpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>> + 'b,
@@ -3625,11 +3625,11 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + 'static>
 
     fn transport_mles<'a>(
         &self,
-        mles: &[MultilinearExtension<'a, E>],
+        mles: Vec<MultilinearExtension<'a, E>>,
     ) -> Vec<Arc<<GpuBackend<E, PCS> as ProverBackend>::MultilinearPoly<'a>>> {
         let cuda_hal = get_cuda_hal().unwrap();
-        mles.iter()
-            .map(|mle| Arc::new(MultilinearExtensionGpu::from_ceno(&cuda_hal, mle)))
+        mles.into_iter()
+            .map(|mle| Arc::new(MultilinearExtensionGpu::from_ceno(&cuda_hal, &mle)))
             .collect_vec()
     }
 }

ceno_zkvm/src/scheme/hal.rs

@@ -117,15 +117,15 @@ pub trait TraceCommitter<PB: ProverBackend> {
         &self,
         traces: BTreeMap<usize, witness::RowMajorMatrix<<PB::E as ExtensionField>::BaseField>>,
     ) -> (
-        Vec<PB::MultilinearPoly<'a>>,
+        Vec<Arc<PB::MultilinearPoly<'a>>>,
         PB::PcsData,
         <PB::Pcs as PolynomialCommitmentScheme<PB::E>>::Commitment,
     );
 
     /// Return an iterator over witness polynomials so backends can decide how to source them
     fn extract_witness_mles<'a, 'b>(
         &self,
-        witness_mles: &'b mut Vec<PB::MultilinearPoly<'a>>,
+        witness_mles: &'b mut Vec<Arc<PB::MultilinearPoly<'a>>>,
         pcs_data: &'b PB::PcsData, // used by GPU backend
     ) -> Box<dyn Iterator<Item = Arc<PB::MultilinearPoly<'a>>> + 'b>;
 }
@@ -280,7 +280,7 @@ pub trait DeviceTransporter<PB: ProverBackend> {
 
     fn transport_mles<'a>(
         &self,
-        mles: &[MultilinearExtension<'a, PB::E>],
+        mles: Vec<MultilinearExtension<'a, PB::E>>,
     ) -> Vec<Arc<PB::MultilinearPoly<'a>>>;
 }
 

ceno_zkvm/src/scheme/prover.rs

@@ -476,7 +476,7 @@ impl<
             // commit to witness traces in batch
             #[cfg_attr(not(feature = "gpu"), allow(unused_mut))]
             let (witness_mles, witness_data, witin_commit): (
-                Vec<PB::MultilinearPoly<'_>>,
+                Vec<Arc<PB::MultilinearPoly<'_>>>,
                 PB::PcsData,
                 PCS::Commitment,
             ) = {
@@ -490,7 +490,8 @@ impl<
                                 gpu_device,
                                 gpu_witness_traces,
                             );
-                        let witness_mles = unsafe { std::mem::transmute(gpu_witness_mles) };
+                        drop(gpu_witness_mles);
+                        let witness_mles = Vec::new();
                         let witness_data = unsafe { std::mem::transmute_copy(&gpu_witness_data) };
                         std::mem::forget(gpu_witness_data);
                         (witness_mles, witness_data, witin_commit)
@@ -894,7 +895,7 @@ impl<
         name_and_instances: Vec<(String, [usize; 2])>,
         structural_rmms: Vec<witness::RowMajorMatrix<E::BaseField>>,
         #[cfg(feature = "gpu")] witness_trace_rows: Vec<Option<usize>>,
-        #[allow(unused_mut)] mut witness_mles: Vec<PB::MultilinearPoly<'data>>,
+        #[allow(unused_mut)] mut witness_mles: Vec<Arc<PB::MultilinearPoly<'data>>>,
         witness_data: &PB::PcsData,
         mut fixed_mles: Vec<Arc<PB::MultilinearPoly<'data>>>,
         challenges: [E; 2],
@@ -967,7 +968,7 @@ impl<
                 let structural_witness = info_span!("[ceno] transport_structural_witness")
                     .in_scope(|| {
                         let structural_mles = structural_rmm.to_mles();
-                        self.device.transport_mles(&structural_mles)
+                        self.device.transport_mles(structural_mles)
                     });
                 (witness_mle, structural_witness, None)
             };