Wgpu/Clamp Kernels #866
Wgpu/Clamp Kernels #866
Conversation
burn-candle/src/ops/tensor.rs
Outdated
| fn clamp_min<const D: usize>( | ||
| tensor: FloatTensor<Self, D>, | ||
| min: FloatElem<Self>, | ||
| ) -> FloatTensor<Self, D> { | ||
| let mask = Self::lower_elem(tensor.clone(), min); | ||
| Self::mask_fill(tensor, mask, min) | ||
| } | ||
|
|
||
| fn clamp_max<const D: usize>( | ||
| tensor: FloatTensor<Self, D>, | ||
| max: FloatElem<Self>, | ||
| ) -> FloatTensor<Self, D> { | ||
| let mask = Self::greater_elem(tensor.clone(), max); | ||
| Self::mask_fill(tensor, mask, max) | ||
| } | ||
|
|
||
| fn clamp<const D: usize>( | ||
| tensor: FloatTensor<Self, D>, | ||
| min: FloatElem<Self>, | ||
| max: FloatElem<Self>, | ||
| ) -> FloatTensor<Self, D> { | ||
| Self::clamp_min(Self::clamp_max(tensor, max), min) | ||
| } |
There was a problem hiding this comment.
Those are the default implementations, we should implement those methods only if we actually add custom implementations.
There was a problem hiding this comment.
Should I take those out of here and put them back in burn-tensor/src/tensor/ops/tensor.rs?
There was a problem hiding this comment.
The main problem i see is that your kernel is always in-place. Normally we like to have a kernel for when tensors can mut (in-place), and another for when they cannot (because there are too many references to them), which will create an output buffer. This logic is already implemented, look for unary ops. Using the unary template, you will only have to provide the logic for each value. For an example you can look at relu in burn-wgpu/src/ops/activation_ops.rs
Also, creating a third kernel for clamp (with both lower and upper bounds) would be a very low hanging fruit once you have made the other two, and it would be a great optimization as we would only launch one kernel instead of two in this case.
|
@louisfd I see what you're saying. I'm not super familiar with proc macros, so correct me if I'm wrong, but don't they expect static literal inputs? In other words, the |
|
@agelas |
|
@louisfd I think this is closer to what you're looking for now. There's no proc macros that really fit the pattern I'd need to make use of wgsl's built-in |
There was a problem hiding this comment.
It looks very good now!
For candle, if we only implement clamp and not clamp_min and clamp_max, then those two will fall back to their default implementation that uses a combination of kernels, so they will be slower than clamp. Looking at candle's code, it seems clamp is just a combination of minimum and maximum which are public, so it should be trivial to have clamp_min call maximum and clamp_max call minimum
|
@louisfd Done! |
Pull Request Template
Checklist
run-checksscript has been executed.Related Issues/PRs
#549
Changes
clampkernels and shaders to wgpu backend.mask_fillmethod that used to be the former default for the candle backend. I took a peak atcandle-coreand it looks like they added their ownclampmethod about a month ago, so we can update onceburn-candleis pinned to a more recent version.Testing
Linked backends to their respective
clamp_min/clamp_maximplementations instead ofmask_fill.