1use std::collections::HashMap;
7
8use crate::gpu::kernels::{
9 BaseKernel, DataType, GpuKernel, KernelMetadata, KernelParams, OperationType,
10};
11use crate::gpu::{GpuBackend, GpuError};
12
13pub struct ElementwiseAddKernel {
15 base: BaseKernel,
16}
17
18impl Default for ElementwiseAddKernel {
19 fn default() -> Self {
20 Self::new()
21 }
22}
23
24impl ElementwiseAddKernel {
25 pub fn new() -> Self {
27 let metadata = KernelMetadata {
28 workgroup_size: [256, 1, 1],
29 local_memory_usage: 0,
30 supports_tensor_cores: false,
31 operationtype: OperationType::MemoryIntensive,
32 backend_metadata: HashMap::new(),
33 };
34
35 let cuda_source = r#"
36extern "C" __global__ void elementwise_add(
37 const float* __restrict__ a,
38 const float* __restrict__ b,
39 float* __restrict__ result,
40 int n
41) {
42 int i = blockIdx.x * blockDim.x + threadIdx.x;
43 if (i < n) {
44 result[i] = a[i] + b[i];
45 }
46}
47"#
48 .to_string();
49
50 let rocm_source = cuda_source.clone();
51
52 let wgpu_source = r#"
53struct Uniforms {
54 n: u32,
55};
56
57@group(0) @binding(0) var<uniform> uniforms: Uniforms;
58@group(0) @binding(1) var<storage, read> a: array<f32>;
59@group(0) @binding(2) var<storage, read> b: array<f32>;
60@group(0) @binding(3) var<storage, write> result: array<f32>;
61
62@compute @workgroup_size(256)
63fn elementwise_add(@builtin(global_invocation_id) global_id: vec3<u32>) {
64 let i = global_id.x;
65
66 if (i < uniforms.n) {
67 result[i] = a[i] + b[i];
68 }
69}
70"#
71 .to_string();
72
73 let metal_source = r#"
74#include <metal_stdlib>
75using namespace metal;
76
77kernel void elementwise_add(
78 const device float* a [[buffer(0)]],
79 const device float* b [[buffer(1)]],
80 device float* result [[buffer(2)]],
81 constant uint& n [[buffer(3)]],
82 uint gid [[thread_position_in_grid]])
83{
84 if (gid < n) {
85 result[gid] = a[gid] + b[gid];
86 }
87}
88"#
89 .to_string();
90
91 let opencl_source = r#"
92__kernel void elementwise_add(
93 __global const float* a,
94 __global const float* b,
95 __global float* result,
96 const int n)
97{
98 int i = get_global_id(0);
99 if (i < n) {
100 result[i] = a[i] + b[i];
101 }
102}
103"#
104 .to_string();
105
106 Self {
107 base: BaseKernel::new(
108 "elementwise_add",
109 &cuda_source,
110 &rocm_source,
111 &wgpu_source,
112 &metal_source,
113 &opencl_source,
114 metadata,
115 ),
116 }
117 }
118}
119
120impl GpuKernel for ElementwiseAddKernel {
121 fn name(&self) -> &str {
122 self.base.name()
123 }
124
125 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
126 self.base.source_for_backend(backend)
127 }
128
129 fn metadata(&self) -> KernelMetadata {
130 self.base.metadata()
131 }
132
133 fn can_specialize(&self, params: &KernelParams) -> bool {
134 matches!(
135 params.datatype,
136 DataType::Float32 | DataType::Float64 | DataType::Int32
137 )
138 }
139
140 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
141 if !self.can_specialize(params) {
142 return Err(GpuError::SpecializationNotSupported);
143 }
144
145 Ok(Box::new(Self::new()))
147 }
148}
149
150pub struct ElementwiseMulKernel {
152 base: BaseKernel,
153}
154
155impl Default for ElementwiseMulKernel {
156 fn default() -> Self {
157 Self::new()
158 }
159}
160
161impl ElementwiseMulKernel {
162 pub fn new() -> Self {
164 let metadata = KernelMetadata {
165 workgroup_size: [256, 1, 1],
166 local_memory_usage: 0,
167 supports_tensor_cores: false,
168 operationtype: OperationType::MemoryIntensive,
169 backend_metadata: HashMap::new(),
170 };
171
172 let cuda_source = r#"
173extern "C" __global__ void elementwise_mul(
174 const float* __restrict__ a,
175 const float* __restrict__ b,
176 float* __restrict__ result,
177 int n
178) {
179 int i = blockIdx.x * blockDim.x + threadIdx.x;
180 if (i < n) {
181 result[i] = a[i] * b[i];
182 }
183}
184"#
185 .to_string();
186
187 let rocm_source = cuda_source.clone();
188
189 let wgpu_source = r#"
190struct Uniforms {
191 n: u32,
192};
193
194@group(0) @binding(0) var<uniform> uniforms: Uniforms;
195@group(0) @binding(1) var<storage, read> a: array<f32>;
196@group(0) @binding(2) var<storage, read> b: array<f32>;
197@group(0) @binding(3) var<storage, write> result: array<f32>;
198
199@compute @workgroup_size(256)
200fn elementwise_mul(@builtin(global_invocation_id) global_id: vec3<u32>) {
201 let i = global_id.x;
202
203 if (i < uniforms.n) {
204 result[i] = a[i] * b[i];
205 }
206}
207"#
208 .to_string();
209
210 let metal_source = r#"
211#include <metal_stdlib>
212using namespace metal;
213
214kernel void elementwise_mul(
215 const device float* a [[buffer(0)]],
216 const device float* b [[buffer(1)]],
217 device float* result [[buffer(2)]],
218 constant uint& n [[buffer(3)]],
219 uint gid [[thread_position_in_grid]])
220{
221 if (gid < n) {
222 result[gid] = a[gid] * b[gid];
223 }
224}
225"#
226 .to_string();
227
228 let opencl_source = r#"
229__kernel void elementwise_mul(
230 __global const float* a,
231 __global const float* b,
232 __global float* result,
233 const int n)
234{
235 int i = get_global_id(0);
236 if (i < n) {
237 result[i] = a[i] * b[i];
238 }
239}
240"#
241 .to_string();
242
243 Self {
244 base: BaseKernel::new(
245 "elementwise_mul",
246 &cuda_source,
247 &rocm_source,
248 &wgpu_source,
249 &metal_source,
250 &opencl_source,
251 metadata,
252 ),
253 }
254 }
255}
256
257impl GpuKernel for ElementwiseMulKernel {
258 fn name(&self) -> &str {
259 self.base.name()
260 }
261
262 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
263 self.base.source_for_backend(backend)
264 }
265
266 fn metadata(&self) -> KernelMetadata {
267 self.base.metadata()
268 }
269
270 fn can_specialize(&self, params: &KernelParams) -> bool {
271 matches!(
272 params.datatype,
273 DataType::Float32 | DataType::Float64 | DataType::Int32
274 )
275 }
276
277 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
278 if !self.can_specialize(params) {
279 return Err(GpuError::SpecializationNotSupported);
280 }
281
282 Ok(Box::new(Self::new()))
283 }
284}
285
286pub struct ElementwiseDivKernel {
288 base: BaseKernel,
289}
290
291impl Default for ElementwiseDivKernel {
292 fn default() -> Self {
293 Self::new()
294 }
295}
296
297impl ElementwiseDivKernel {
298 pub fn new() -> Self {
300 let metadata = KernelMetadata {
301 workgroup_size: [256, 1, 1],
302 local_memory_usage: 0,
303 supports_tensor_cores: false,
304 operationtype: OperationType::ComputeIntensive, backend_metadata: HashMap::new(),
306 };
307
308 let cuda_source = r#"
309extern "C" __global__ void elementwise_div(
310 const float* __restrict__ a,
311 const float* __restrict__ b,
312 float* __restrict__ result,
313 int n
314) {
315 int i = blockIdx.x * blockDim.x + threadIdx.x;
316 if (i < n) {
317 result[i] = a[i] / b[i];
318 }
319}
320"#
321 .to_string();
322
323 let rocm_source = cuda_source.clone();
324
325 let wgpu_source = r#"
326struct Uniforms {
327 n: u32,
328};
329
330@group(0) @binding(0) var<uniform> uniforms: Uniforms;
331@group(0) @binding(1) var<storage, read> a: array<f32>;
332@group(0) @binding(2) var<storage, read> b: array<f32>;
333@group(0) @binding(3) var<storage, write> result: array<f32>;
334
335@compute @workgroup_size(256)
336fn elementwise_div(@builtin(global_invocation_id) global_id: vec3<u32>) {
337 let i = global_id.x;
338
339 if (i < uniforms.n) {
340 result[i] = a[i] / b[i];
341 }
342}
343"#
344 .to_string();
345
346 let metal_source = r#"
347#include <metal_stdlib>
348using namespace metal;
349
350kernel void elementwise_div(
351 const device float* a [[buffer(0)]],
352 const device float* b [[buffer(1)]],
353 device float* result [[buffer(2)]],
354 constant uint& n [[buffer(3)]],
355 uint gid [[thread_position_in_grid]])
356{
357 if (gid < n) {
358 result[gid] = a[gid] / b[gid];
359 }
360}
361"#
362 .to_string();
363
364 let opencl_source = r#"
365__kernel void elementwise_div(
366 __global const float* a,
367 __global const float* b,
368 __global float* result,
369 const int n)
370{
371 int i = get_global_id(0);
372 if (i < n) {
373 result[i] = a[i] / b[i];
374 }
375}
376"#
377 .to_string();
378
379 Self {
380 base: BaseKernel::new(
381 "elementwise_div",
382 &cuda_source,
383 &rocm_source,
384 &wgpu_source,
385 &metal_source,
386 &opencl_source,
387 metadata,
388 ),
389 }
390 }
391}
392
393impl GpuKernel for ElementwiseDivKernel {
394 fn name(&self) -> &str {
395 self.base.name()
396 }
397
398 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
399 self.base.source_for_backend(backend)
400 }
401
402 fn metadata(&self) -> KernelMetadata {
403 self.base.metadata()
404 }
405
406 fn can_specialize(&self, params: &KernelParams) -> bool {
407 matches!(params.datatype, DataType::Float32 | DataType::Float64)
408 }
409
410 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
411 if !self.can_specialize(params) {
412 return Err(GpuError::SpecializationNotSupported);
413 }
414
415 Ok(Box::new(Self::new()))
416 }
417}
418
419pub struct ScalarMulKernel {
421 base: BaseKernel,
422}
423
424impl Default for ScalarMulKernel {
425 fn default() -> Self {
426 Self::new()
427 }
428}
429
430impl ScalarMulKernel {
431 pub fn new() -> Self {
433 let metadata = KernelMetadata {
434 workgroup_size: [256, 1, 1],
435 local_memory_usage: 0,
436 supports_tensor_cores: false,
437 operationtype: OperationType::MemoryIntensive,
438 backend_metadata: HashMap::new(),
439 };
440
441 let cuda_source = r#"
442extern "C" __global__ void scalar_mul(
443 const float* __restrict__ input,
444 float* __restrict__ output,
445 float scalar,
446 int n
447) {
448 int i = blockIdx.x * blockDim.x + threadIdx.x;
449 if (i < n) {
450 output[i] = input[i] * scalar;
451 }
452}
453"#
454 .to_string();
455
456 let rocm_source = cuda_source.clone();
457
458 let wgpu_source = r#"
459struct Uniforms {
460 n: u32,
461 scalar: f32,
462};
463
464@group(0) @binding(0) var<uniform> uniforms: Uniforms;
465@group(0) @binding(1) var<storage, read> input: array<f32>;
466@group(0) @binding(2) var<storage, write> output: array<f32>;
467
468@compute @workgroup_size(256)
469fn scalar_mul(@builtin(global_invocation_id) global_id: vec3<u32>) {
470 let i = global_id.x;
471
472 if (i < uniforms.n) {
473 output[i] = input[i] * uniforms.scalar;
474 }
475}
476"#
477 .to_string();
478
479 let metal_source = r#"
480#include <metal_stdlib>
481using namespace metal;
482
483kernel void scalar_mul(
484 const device float* input [[buffer(0)]],
485 device float* output [[buffer(1)]],
486 constant float& scalar [[buffer(2)]],
487 constant uint& n [[buffer(3)]],
488 uint gid [[thread_position_in_grid]])
489{
490 if (gid < n) {
491 output[gid] = input[gid] * scalar;
492 }
493}
494"#
495 .to_string();
496
497 let opencl_source = r#"
498__kernel void scalar_mul(
499 __global const float* input,
500 __global float* output,
501 const float scalar,
502 const int n)
503{
504 int i = get_global_id(0);
505 if (i < n) {
506 output[i] = input[i] * scalar;
507 }
508}
509"#
510 .to_string();
511
512 Self {
513 base: BaseKernel::new(
514 "scalar_mul",
515 &cuda_source,
516 &rocm_source,
517 &wgpu_source,
518 &metal_source,
519 &opencl_source,
520 metadata,
521 ),
522 }
523 }
524}
525
526impl GpuKernel for ScalarMulKernel {
527 fn name(&self) -> &str {
528 self.base.name()
529 }
530
531 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
532 self.base.source_for_backend(backend)
533 }
534
535 fn metadata(&self) -> KernelMetadata {
536 self.base.metadata()
537 }
538
539 fn can_specialize(&self, params: &KernelParams) -> bool {
540 matches!(
541 params.datatype,
542 DataType::Float32 | DataType::Float64 | DataType::Int32
543 )
544 }
545
546 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
547 if !self.can_specialize(params) {
548 return Err(GpuError::SpecializationNotSupported);
549 }
550
551 Ok(Box::new(Self::new()))
552 }
553}
554pub struct ElementwiseSubKernel {
556 base: BaseKernel,
557}
558
559impl Default for ElementwiseSubKernel {
560 fn default() -> Self {
561 Self::new()
562 }
563}
564
565impl ElementwiseSubKernel {
566 pub fn new() -> Self {
567 let metadata = KernelMetadata {
568 workgroup_size: [256, 1, 1],
569 local_memory_usage: 0,
570 supports_tensor_cores: false,
571 operationtype: OperationType::MemoryIntensive,
572 backend_metadata: HashMap::new(),
573 };
574
575 let cuda_source = r#"
576extern \"C\" __global__ void elementwise_sub(
577 const float* __restrict__ a,
578 const float* __restrict__ b,
579 float* __restrict__ result,
580 int n
581) {
582 int i = blockIdx.x * blockDim.x + threadIdx.x;
583 if (i < n) {
584 result[i] = a[i] - b[i];
585 }
586}
587"#
588 .to_string();
589
590 let rocm_source = cuda_source.clone();
591
592 let wgpu_source = r#"
593struct Uniforms {
594 n: u32,
595};
596
597@group(0) @binding(0) var<uniform> uniforms: Uniforms;
598@group(0) @binding(1) var<storage, read> a: array<f32>;
599@group(0) @binding(2) var<storage, read> b: array<f32>;
600@group(0) @binding(3) var<storage, read_write> result: array<f32>;
601
602@compute @workgroup_size(256)
603fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
604 let i = global_id.x;
605
606 if (i < uniforms.n) {
607 result[i] = a[i] - b[i];
608 }
609}
610"#
611 .to_string();
612
613 let metal_source = String::new();
614 let opencl_source = String::new();
615
616 Self {
617 base: BaseKernel::new(
618 "elementwise_sub",
619 &cuda_source,
620 &rocm_source,
621 &wgpu_source,
622 &metal_source,
623 &opencl_source,
624 metadata,
625 ),
626 }
627 }
628}
629
630impl GpuKernel for ElementwiseSubKernel {
631 fn name(&self) -> &str {
632 self.base.name()
633 }
634
635 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
636 self.base.source_for_backend(backend)
637 }
638
639 fn metadata(&self) -> KernelMetadata {
640 self.base.metadata()
641 }
642
643 fn can_specialize(&self, params: &KernelParams) -> bool {
644 matches!(
645 params.datatype,
646 DataType::Float32 | DataType::Float64 | DataType::Int32
647 )
648 }
649
650 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
651 if !self.can_specialize(params) {
652 return Err(GpuError::SpecializationNotSupported);
653 }
654 Ok(Box::new(Self::new()))
655 }
656}
657pub struct ElementwisePowKernel {
659 base: BaseKernel,
660}
661
662impl Default for ElementwisePowKernel {
663 fn default() -> Self {
664 Self::new()
665 }
666}
667
668impl ElementwisePowKernel {
669 pub fn new() -> Self {
670 let metadata = KernelMetadata {
671 workgroup_size: [256, 1, 1],
672 local_memory_usage: 0,
673 supports_tensor_cores: false,
674 operationtype: OperationType::ComputeIntensive,
675 backend_metadata: HashMap::new(),
676 };
677
678 let cuda_source = r#"
679extern "C" __global__ void elementwise_pow(
680 const float* __restrict__ a,
681 const float* __restrict__ b,
682 float* __restrict__ result,
683 int n
684) {
685 int i = blockIdx.x * blockDim.x + threadIdx.x;
686 if (i < n) {
687 result[i] = powf(a[i], b[i]);
688 }
689}
690"#
691 .to_string();
692
693 let rocm_source = cuda_source.clone();
694
695 let wgpu_source = r#"
696struct Uniforms {
697 n: u32,
698};
699
700@group(0) @binding(0) var<uniform> uniforms: Uniforms;
701@group(0) @binding(1) var<storage, read> a: array<f32>;
702@group(0) @binding(2) var<storage, read> b: array<f32>;
703@group(0) @binding(3) var<storage, read_write> result: array<f32>;
704
705@compute @workgroup_size(256)
706fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
707 let i = global_id.x;
708
709 if (i < uniforms.n) {
710 result[i] = pow(a[i], b[i]);
711 }
712}
713"#
714 .to_string();
715
716 let metal_source = String::new();
717 let opencl_source = String::new();
718
719 Self {
720 base: BaseKernel::new(
721 "elementwise_pow",
722 &cuda_source,
723 &rocm_source,
724 &wgpu_source,
725 &metal_source,
726 &opencl_source,
727 metadata,
728 ),
729 }
730 }
731}
732
733impl GpuKernel for ElementwisePowKernel {
734 fn name(&self) -> &str {
735 self.base.name()
736 }
737
738 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
739 self.base.source_for_backend(backend)
740 }
741
742 fn metadata(&self) -> KernelMetadata {
743 self.base.metadata()
744 }
745
746 fn can_specialize(&self, params: &KernelParams) -> bool {
747 matches!(params.datatype, DataType::Float32 | DataType::Float64)
748 }
749
750 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
751 if !self.can_specialize(params) {
752 return Err(GpuError::SpecializationNotSupported);
753 }
754 Ok(Box::new(Self::new()))
755 }
756}
757
758pub struct ElementwiseSqrtKernel {
760 base: BaseKernel,
761}
762
763impl Default for ElementwiseSqrtKernel {
764 fn default() -> Self {
765 Self::new()
766 }
767}
768
769impl ElementwiseSqrtKernel {
770 pub fn new() -> Self {
771 let metadata = KernelMetadata {
772 workgroup_size: [256, 1, 1],
773 local_memory_usage: 0,
774 supports_tensor_cores: false,
775 operationtype: OperationType::ComputeIntensive,
776 backend_metadata: HashMap::new(),
777 };
778
779 let cuda_source = r#"
780extern "C" __global__ void elementwise_sqrt(
781 const float* __restrict__ a,
782 float* __restrict__ result,
783 int n
784) {
785 int i = blockIdx.x * blockDim.x + threadIdx.x;
786 if (i < n) {
787 result[i] = sqrtf(a[i]);
788 }
789}
790"#
791 .to_string();
792
793 let rocm_source = cuda_source.clone();
794
795 let wgpu_source = r#"
796struct Uniforms {
797 n: u32,
798};
799
800@group(0) @binding(0) var<uniform> uniforms: Uniforms;
801@group(0) @binding(1) var<storage, read> input: array<f32>;
802@group(0) @binding(2) var<storage, read_write> result: array<f32>;
803
804@compute @workgroup_size(256)
805fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
806 let i = global_id.x;
807
808 if (i < uniforms.n) {
809 result[i] = sqrt(input[i]);
810 }
811}
812"#
813 .to_string();
814
815 let metal_source = String::new();
816 let opencl_source = String::new();
817
818 Self {
819 base: BaseKernel::new(
820 "elementwise_sqrt",
821 &cuda_source,
822 &rocm_source,
823 &wgpu_source,
824 &metal_source,
825 &opencl_source,
826 metadata,
827 ),
828 }
829 }
830}
831
832impl GpuKernel for ElementwiseSqrtKernel {
833 fn name(&self) -> &str {
834 self.base.name()
835 }
836
837 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
838 self.base.source_for_backend(backend)
839 }
840
841 fn metadata(&self) -> KernelMetadata {
842 self.base.metadata()
843 }
844
845 fn can_specialize(&self, params: &KernelParams) -> bool {
846 matches!(params.datatype, DataType::Float32 | DataType::Float64)
847 }
848
849 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
850 if !self.can_specialize(params) {
851 return Err(GpuError::SpecializationNotSupported);
852 }
853 Ok(Box::new(Self::new()))
854 }
855}
856
857pub struct ElementwiseExpKernel {
859 base: BaseKernel,
860}
861
862impl Default for ElementwiseExpKernel {
863 fn default() -> Self {
864 Self::new()
865 }
866}
867
868impl ElementwiseExpKernel {
869 pub fn new() -> Self {
870 let metadata = KernelMetadata {
871 workgroup_size: [256, 1, 1],
872 local_memory_usage: 0,
873 supports_tensor_cores: false,
874 operationtype: OperationType::ComputeIntensive,
875 backend_metadata: HashMap::new(),
876 };
877
878 let cuda_source = r#"
879extern "C" __global__ void elementwise_exp(
880 const float* __restrict__ a,
881 float* __restrict__ result,
882 int n
883) {
884 int i = blockIdx.x * blockDim.x + threadIdx.x;
885 if (i < n) {
886 result[i] = expf(a[i]);
887 }
888}
889"#
890 .to_string();
891
892 let rocm_source = cuda_source.clone();
893
894 let wgpu_source = r#"
895struct Uniforms {
896 n: u32,
897};
898
899@group(0) @binding(0) var<uniform> uniforms: Uniforms;
900@group(0) @binding(1) var<storage, read> input: array<f32>;
901@group(0) @binding(2) var<storage, read_write> result: array<f32>;
902
903@compute @workgroup_size(256)
904fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
905 let i = global_id.x;
906
907 if (i < uniforms.n) {
908 result[i] = exp(input[i]);
909 }
910}
911"#
912 .to_string();
913
914 let metal_source = String::new();
915 let opencl_source = String::new();
916
917 Self {
918 base: BaseKernel::new(
919 "elementwise_exp",
920 &cuda_source,
921 &rocm_source,
922 &wgpu_source,
923 &metal_source,
924 &opencl_source,
925 metadata,
926 ),
927 }
928 }
929}
930
931impl GpuKernel for ElementwiseExpKernel {
932 fn name(&self) -> &str {
933 self.base.name()
934 }
935
936 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
937 self.base.source_for_backend(backend)
938 }
939
940 fn metadata(&self) -> KernelMetadata {
941 self.base.metadata()
942 }
943
944 fn can_specialize(&self, params: &KernelParams) -> bool {
945 matches!(params.datatype, DataType::Float32 | DataType::Float64)
946 }
947
948 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
949 if !self.can_specialize(params) {
950 return Err(GpuError::SpecializationNotSupported);
951 }
952 Ok(Box::new(Self::new()))
953 }
954}
955
956pub struct ElementwiseLogKernel {
958 base: BaseKernel,
959}
960
961impl Default for ElementwiseLogKernel {
962 fn default() -> Self {
963 Self::new()
964 }
965}
966
967impl ElementwiseLogKernel {
968 pub fn new() -> Self {
969 let metadata = KernelMetadata {
970 workgroup_size: [256, 1, 1],
971 local_memory_usage: 0,
972 supports_tensor_cores: false,
973 operationtype: OperationType::ComputeIntensive,
974 backend_metadata: HashMap::new(),
975 };
976
977 let cuda_source = r#"
978extern "C" __global__ void elementwise_log(
979 const float* __restrict__ a,
980 float* __restrict__ result,
981 int n
982) {
983 int i = blockIdx.x * blockDim.x + threadIdx.x;
984 if (i < n) {
985 result[i] = logf(a[i]);
986 }
987}
988"#
989 .to_string();
990
991 let rocm_source = cuda_source.clone();
992
993 let wgpu_source = r#"
994struct Uniforms {
995 n: u32,
996};
997
998@group(0) @binding(0) var<uniform> uniforms: Uniforms;
999@group(0) @binding(1) var<storage, read> input: array<f32>;
1000@group(0) @binding(2) var<storage, read_write> result: array<f32>;
1001
1002@compute @workgroup_size(256)
1003fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
1004 let i = global_id.x;
1005
1006 if (i < uniforms.n) {
1007 result[i] = log(input[i]);
1008 }
1009}
1010"#
1011 .to_string();
1012
1013 let metal_source = String::new();
1014 let opencl_source = String::new();
1015
1016 Self {
1017 base: BaseKernel::new(
1018 "elementwise_log",
1019 &cuda_source,
1020 &rocm_source,
1021 &wgpu_source,
1022 &metal_source,
1023 &opencl_source,
1024 metadata,
1025 ),
1026 }
1027 }
1028}
1029
1030impl GpuKernel for ElementwiseLogKernel {
1031 fn name(&self) -> &str {
1032 self.base.name()
1033 }
1034
1035 fn source_for_backend(&self, backend: GpuBackend) -> Result<String, GpuError> {
1036 self.base.source_for_backend(backend)
1037 }
1038
1039 fn metadata(&self) -> KernelMetadata {
1040 self.base.metadata()
1041 }
1042
1043 fn can_specialize(&self, params: &KernelParams) -> bool {
1044 matches!(params.datatype, DataType::Float32 | DataType::Float64)
1045 }
1046
1047 fn specialize(&self, params: &KernelParams) -> Result<Box<dyn GpuKernel>, GpuError> {
1048 if !self.can_specialize(params) {
1049 return Err(GpuError::SpecializationNotSupported);
1050 }
1051 Ok(Box::new(Self::new()))
1052 }
1053}