and_broadcast_mut for Zip

[millardjn]

Currently Zip has no mutable equivalent for the and_broadcast method, which is presumably to allow it to safely be split() and parallelized. However this makes some calculations difficult. The examples below are for backprop through a broadcast-able element-wise multiplication, but the problem appears as a pattern often.

Currently I'm looking at using this:

unsafe{
	let input2_grad = data.get_mut(&self.input2_id.gradient_id())?;

	// do not split/parallelise this Zip!
	Zip::from(&input1)
		.and(&output_grad)
		.and_broadcast(&input2_grad)
		.apply(|input1, out_grad, input2_grad| {
			let input2_grad = input2_grad as *const f32 as *mut f32;
			*input2_grad += input1 * out_grad;
		});
}

which works correctly, but also has UB written all over it depending on how Zip is implemented internally (I realize Zip owes me nothing here).

The safe alternative I've found is:

let iter = input1.exact_chunks(input2.shape()).into_iter()
	.zip(output_grad.exact_chunks(input2.shape()));

for (input1_chunk, out_grad_chunk) in iter {
	Zip::from(&mut input2_grad)
		.and(&input1_chunk)
		.and(&out_grad_chunk)
		.apply(|input2_grad, input1, out_grad| {
			*input2_grad += input1 * out_grad;
		});
}

which is catastrophically (over 30x) slower for some cases.

So my questions are:

• is there a safe high performance workaround that I've missed?
• would it be possible to add a and_broadcast_mut() method which would set a flag or something which makes split() panic?

[bluss]

Broadcasting isn't compatible with mutable access in general. We can use a Vec or 1D array as a model for this. We assume throughout the code that if the element index is different, two mutable references to elements are not aliasing, so we can modify though both paths. This breaks if we have mutable broadcasting.

It could only be allowed if we imposed something like the access pattern of a Cell, and not allowing mutable references.

Which brings a "solution": If an ArrayViewMut<T> can be converted into an ArrayView<Cell<T>> elements, which is a transformation that has been discussed before and we have a similar example here in libstd https://doc.rust-lang.org/nightly/std/cell/struct.Cell.html#method.as_slice_of_cells.

[bluss]

In short: .and_broadcast_cell ? Or General ArrayViewMut -> ArrayView of Cells?

One cool fact is that the non-parallelization should work itself out automatically by the Send/Sync traits; when the Zip contains a Cell, the compiler will disallow thread parallelism.

[millardjn]

I think ArrayViewMut<T> borrowed as a ArrayView<Cell<T>> looks like quite a nice solution, and might even come in handy elsewhere. From this it seems simple, are there any pitfalls in extending this to casting the backing array? I'll give implementing this a go.

[bluss]

We should be able to start with just an ArrayViewMut::as_cell_view, or whatever name feels best, method and try to plug it in the existing broadcast code. Ergonomics of Cells and value updates are of course not ideal.

I guess there is a chance these methods are being held back for technical or "unsafe code guidelines/formalization" reasons upstream. I can't spot any discussion along those lines though. The Cell has this attribute in Rust now, which is a stable attribute too: #[repr(transparent)].

[bluss]

. o O ( So what if there's an arithmetic compatible transparent wrapper type around Cell )

[jturner314]

The Cell has this attribute in Rust now, which is a stable attribute too: #[repr(transparent)].

#[repr(transparent)] (at least on Cell) is in beta, but not stable yet. (It should be released as stable in a couple of weeks.) (stable docs, beta docs)

So what if there's an arithmetic compatible transparent wrapper type around Cell

I think the #[repr(transparent)] annotation is technically relevant for our use-case because AFAIK Rust doesn't guarantee any particular bitwise representation by default except for fixed-size arrays. (So Cell needs to have #[repr(C)] or #[repr(transparent)] for us to be able to safely cast the ArrayViewMut in-place.) In practice, though, I don't think we have to worry too much about this because Cell<T> is just a thin wrapper around UnsafeCell<T> which is just a thin wrapper around T. Regardless, it'll be resolved in a couple of weeks anyway.

[bluss]

I wonder, after we start supporting Cell there will be a want for supporting something that wraps Cell (with the same features) that also supports arithmetic. So that it can participate in operations like Array1<f32> + ArrayView1<MathCell<f32>>