The most effective way to transform each lane for generic dimensions.

[Wleter]

Hi, appreciate this crate it's really well thought. I am using it in scientific project where performance is important, so I have a question
whether I implemented transformations of each lane in a generic ndarray the best way possible or not.

I am wondering because when I looked at flamegraph of my calculations and almost 70% of my time is spend on heap allocation and freeing, after adding those operations. I tried many changes but I didn't succeed.

The first transformation is like a generalized matrix multiplying (I am also wondering whether there is a way to multiply not 1d lane but 2d "multi lane" if it is possible)

fn matrix_transform(&mut self, array: &mut Array<Complex64, N>) {
    array.lanes_mut(Axis(self.dimension_no))
        .into_iter()
        .par_bridge()
        .for_each(|mut lane| lane.assign(&self.transformation.dot(&lane)));
}

and the second transformation transform each lane using fft

fn fft_transform(&mut self, array: &mut Array<Complex64, N>) {
    let dimension_size_sqrt = (self.dimension_size as f64).sqrt();

    array.lanes_mut(Axis(self.dimension_no))
        .into_iter()
        .par_bridge()
        .for_each(|mut lane| {
            let mut temp = lane.to_vec();
            self.fft.process(&mut temp);

            lane.iter_mut().zip(temp.iter()).for_each(|(dest, src)| {
                *dest = *src / dimension_size_sqrt;
            });
        });
}

Thanks in advance!

[bluss]

I haven't looked into the details, but I'd recommend first trying:

1. Don't use iterators and std::iter::zip
2. Use ndarray::Zip and its parallelization support
3. Don't use par_bridge, it is slow, unless you have to

It looks like you need a function like general_mat_vec_mul which ndarray has right there, but maybe I read the code wrong.

[Wleter]

Thanks, because of python I only think about zip when I zip two things and I didn't thought about Zip for this case.

Raw speed and allocation time is the same but parallel is now 8% faster. As I understand the assign method copies the result to
lane, is it possible to move the lane to result instead of copying the result? I was focused on this part and didn't succeed.

Now my code looks like this:

Zip::from(array.lanes_mut(Axis(self.transform_dim)))
    .par_for_each(|mut lane| lane.assign(&self.transformation.dot(&lane)));

and

Zip::from(array.lanes_mut(Axis(self.transform_dim))).par_for_each(
    |mut lane| {
        let mut temp = lane.to_vec();
        self.fft.process(&mut temp);

        lane.iter_mut().zip(temp.iter()).for_each(|(dest, src)| {
            *dest = *src / dimension_size_sqrt;
        });
    },
)

I don't think general_mat_vec_mul can be used here since i would need to borrow two times lane:
general_mat_vec_mul(Complex64::one(), self.transformation, &lane, Complex64::zero(), &mut lane) or am I misunderstanding.

[bluss]

You still have one iterator zip in there, replace it with Zip and it will iterate faster (don't need to be parallel for that inner zip I think). 🙂 Zip supports slices and &Vec natively (they are easily converted to arrayviews).

I see, the in place matrix multiply is tricky. Maybe use an intermediate array that you multiply into, instead of allocating a new one for each lane?

As for the first matrix_transform, is it also just something like X' = AX where A, X are matrices?

[nilgoyette]

Depending on the size of temp and the number of iterations, you might want to use "thread data". You can use a struct if you need more than one variable.

// thread_data is cloned only once for each thread, then it's moved within the same thread.
let temp = Array1<Complex64>::zeros(?);

Zip::from(array.lanes_mut(Axis(self.transform_dim))).into_par_iter().for_each_with(
    temp, |mut temp, mut lane| {
        // You can use assign because I used an Array1 instead of a Vec
        temp.assign(&lane);
        // This unwrap can't fail because it's a contiguous 1D array
        self.fft.process(temp.as_slice_mut().unwrap());
        // Put back everything in the lane
        Zip::from(&mut lane, &temp).for_each(|l, t| *l = t / dimension_size_sqrt);
    });

Moreover, you probably want to use process_with_scratch instead of process.

Also, if you're motivated enough. One of the lane will be the actual memory order of your array, so maybe you don't need to copy it in a temporary vector (only in that specific lane). You can use a slice directly from your array.

fft.process_with_scratch(lane.as_slice_memory_order_mut().unwrap(), my_scratch);

[Wleter]

Tanks for replies,

You still have one iterator zip in there, replace it with Zip and it will iterate faster

From benchmarks the iterator seems to be faster than zip.

As for the first matrix_transform, is it also just something like X' = AX where A, X are matrices?

For array that is 2d and for transform_dim = 0 the matrix_transform do this, for 3d and transform_dim = 1,
it does X'_{ijk} = A_{jm} X_{imk} in einstein notation.

Depending on the size of temp and the number of iterations, you might want to use "thread data". You can use a struct if you need more than one variable.

In my case it is also slower than earlier method.
For the process_with_scratch I need to flatten the array accordingly and then reshape it. I don't know a good method to flatten properly.
Also lane.as_slice_memory_order_mut() applied to last dimension is significantly slower than parallel.

It seems that I didn't do that bad in implementing these methods and it is as fast as it could be, apart from not using Zip.

[nilgoyette]

is significantly slower than parallel.

What do you mean by that? It's not an either/or. You can use what I suggested in a parallel Zip.

Btw, if you consider this conversation resolved, please use "Mark as answer".

[Wleter]

I made a mistake, I meant fft.process_with_scratch(lane.as_slice_memory_order_mut().unwrap(), my_scratch); this approach, sorry for confusion.

The one you suggested with parallel Zip and allocating temp array before zipping is just slightly slower.