(self)
| 930 | |
| 931 | @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") |
| 932 | def test_custom_kernel_attributes(self): |
| 933 | if mx.metal.is_available(): |
| 934 | source = "out[0] = threads_per_threadgroup.x;" |
| 935 | custom_kernel = mx.fast.metal_kernel |
| 936 | elif mx.cuda.is_available(): |
| 937 | source = "out[0] = blockDim.x;" |
| 938 | custom_kernel = mx.fast.cuda_kernel |
| 939 | |
| 940 | a = mx.zeros(shape=(1, 1)) |
| 941 | kernel = custom_kernel( |
| 942 | name="test_fun", |
| 943 | input_names=["a"], |
| 944 | output_names=["out"], |
| 945 | source=source, |
| 946 | ) |
| 947 | out = kernel( |
| 948 | inputs=[a], |
| 949 | grid=(2, 1, 1), |
| 950 | threadgroup=(2, 1, 1), |
| 951 | output_shapes=[(1, 1)], |
| 952 | output_dtypes=[mx.uint32], |
| 953 | stream=mx.gpu, |
| 954 | )[0] |
| 955 | self.assertEqual(out.item(), 2) |
| 956 | |
| 957 | @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") |
| 958 | def test_custom_kernel_caching(self): |
nothing calls this directly
no test coverage detected