Repository navigation
and_broadcast_mut for Zip #478
Description
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 makessplit()panic?
Activity
bluss commentedon Aug 26, 2018
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 anArrayView<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 commentedon Aug 26, 2018
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 commentedon Aug 27, 2018
I think
ArrayViewMut<T>borrowed as aArrayView<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 commentedon Aug 27, 2018
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
Cellhas this attribute in Rust now, which is a stable attribute too:#[repr(transparent)].bluss commentedon Aug 27, 2018
. o O ( So what if there's an arithmetic compatible transparent wrapper type around Cell )
jturner314 commentedon Aug 27, 2018
#[repr(transparent)](at least onCell) is in beta, but not stable yet. (It should be released as stable in a couple of weeks.) (stable docs, beta docs)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. (SoCellneeds to have#[repr(C)]or#[repr(transparent)]for us to be able to safely cast theArrayViewMutin-place.) In practice, though, I don't think we have to worry too much about this becauseCell<T>is just a thin wrapper aroundUnsafeCell<T>which is just a thin wrapper aroundT. Regardless, it'll be resolved in a couple of weeks anyway.bluss commentedon Nov 18, 2018
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>>