(self)
| 738 | |
| 739 | @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") |
| 740 | def test_custom_kernel_basic(self): |
| 741 | if mx.metal.is_available(): |
| 742 | source = """ |
| 743 | uint elem = thread_position_in_grid.x; |
| 744 | out1[elem] = a[elem]; |
| 745 | """ |
| 746 | custom_kernel = mx.fast.metal_kernel |
| 747 | elif mx.cuda.is_available(): |
| 748 | source = """ |
| 749 | auto elem = cooperative_groups::this_grid().thread_rank(); |
| 750 | out1[elem] = a[elem]; |
| 751 | """ |
| 752 | custom_kernel = mx.fast.cuda_kernel |
| 753 | |
| 754 | mx.random.seed(7) |
| 755 | a = mx.random.normal(shape=(2, 2)) |
| 756 | kernel = custom_kernel( |
| 757 | name="basic", |
| 758 | input_names=["a"], |
| 759 | output_names=["out1"], |
| 760 | source=source, |
| 761 | ) |
| 762 | out = kernel( |
| 763 | inputs=[a], |
| 764 | grid=(4, 1, 1), |
| 765 | threadgroup=(2, 1, 1), |
| 766 | output_shapes=[(2, 2)], |
| 767 | output_dtypes=[mx.float32], |
| 768 | stream=mx.gpu, |
| 769 | ) |
| 770 | self.assertTrue(mx.allclose(out[0], a)) |
| 771 | |
| 772 | @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") |
| 773 | def test_custom_kernel_args(self): |
nothing calls this directly
no test coverage detected