(
&self,
a: &Tensor,
b: &Tensor,
c: &Tensor,
entry: &str,
)
| 825 | } |
| 826 | |
| 827 | fn dispatch_ternary_vec4( |
| 828 | &self, |
| 829 | a: &Tensor, |
| 830 | b: &Tensor, |
| 831 | c: &Tensor, |
| 832 | entry: &str, |
| 833 | ) -> Result<Tensor> { |
| 834 | let dtype = a.dtype(); |
| 835 | let shape = a.shape().clone(); |
| 836 | let n = a.elem_count(); |
| 837 | |
| 838 | let buf_a = self.get_or_upload(a)?; |
| 839 | let buf_b = self.get_or_upload(b)?; |
| 840 | let buf_c = self.get_or_upload(c)?; |
| 841 | let buf_out = self.alloc_output(n); |
| 842 | |
| 843 | let threads_needed = (n as u32).div_ceil(4); |
| 844 | let result = self.dispatch_compute( |
| 845 | entry, |
| 846 | &[buf_a.buffer, buf_b.buffer, buf_out.buffer, buf_c.buffer], |
| 847 | &buf_out, |
| 848 | n, |
| 849 | &[n as u32, 0, 0, 0], |
| 850 | (threads_needed.div_ceil(WG_ELEM), 1, 1), |
| 851 | ); |
| 852 | |
| 853 | let tensor = Self::from_f32_vec(result, shape.dims(), dtype)?; |
| 854 | self.cache_activation(tensor.id(), buf_out); |
| 855 | Ok(tensor) |
| 856 | } |
| 857 | |
| 858 | fn dispatch_unary_vec4(&self, x: &Tensor, entry: &str) -> Result<Tensor> { |
| 859 | let dtype = x.dtype(); |
no test coverage detected