(self)
| 956 | |
| 957 | @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") |
| 958 | def test_custom_kernel_caching(self): |
| 959 | def call_kernel(a: mx.array, source): |
| 960 | kernel = mx.fast.metal_kernel( |
| 961 | name="my_kernel", |
| 962 | input_names=["inp"], |
| 963 | output_names=["out"], |
| 964 | source=source, |
| 965 | ) |
| 966 | return kernel( |
| 967 | inputs=[a], |
| 968 | grid=(a.size, 1, 1), |
| 969 | threadgroup=(a.size, 1, 1), |
| 970 | output_shapes=[a.shape], |
| 971 | output_dtypes=[a.dtype], |
| 972 | stream=mx.gpu, |
| 973 | )[0] |
| 974 | |
| 975 | a = mx.random.normal(shape=(32,)) |
| 976 | |
| 977 | source = """ |
| 978 | uint elem = thread_position_in_grid.x; |
| 979 | out[elem] = 0.0; |
| 980 | """ |
| 981 | |
| 982 | out = call_kernel(a, source) |
| 983 | self.assertTrue(mx.array_equal(out, mx.zeros_like(out))) |
| 984 | |
| 985 | source = """ |
| 986 | uint elem = thread_position_in_grid.x; |
| 987 | out[elem] = 1.0; |
| 988 | """ |
| 989 | out = call_kernel(a, source) |
| 990 | self.assertTrue(mx.array_equal(out, mx.ones_like(out))) |
| 991 | |
| 992 | |
| 993 | if __name__ == "__main__": |
nothing calls this directly
no outgoing calls
no test coverage detected