Understanding stage vise process in from_values_with_gpu()

[vivekvpandya]

in from_values_with_gpu() I see following code

        if values_num_per_poly*poly_num == ctx.values_flatten.len() {
            println!("in first stage");
            ext_values_flatten = Arc::<Vec<F>>::get_mut(&mut ctx.ext_values_flatten).unwrap();
            values_flatten = Arc::<Vec<F, MyAllocator>>::get_mut(&mut ctx.values_flatten).unwrap();
            digests_and_caps_buf
                = Arc::<Vec<<<C as GenericConfig<D>>::Hasher as Hasher<F>>::Hash>>::get_mut(&mut ctx.digests_and_caps_buf).unwrap();
            ext_values_device_offset = 0;
            is_first_stage = true;
        } else {
        // } else if values_num_per_poly*poly_num == ctx.values_flatten2.len() {
            println!("in second stage");
            ext_values_flatten = Arc::<Vec<F>>::get_mut(&mut ctx.ext_values_flatten2).unwrap();
            values_flatten     = Arc::<Vec<F, MyAllocator>>::get_mut(&mut ctx.values_flatten2).unwrap();
            digests_and_caps_buf
                =  Arc::<Vec<<<C as GenericConfig<D>>::Hasher as Hasher<F>>::Hash>>::get_mut(&mut ctx.digests_and_caps_buf2).unwrap();
            ext_values_device_offset = ctx.second_stage_offset;
            is_first_stage = false;
        }

It seems to be using pre-allocated memory in CTX for different stages. Are this stages some kind of optimisation or its required for correct results?

[vivekvpandya]

@liangping