| 20 | namespace { |
| 21 | |
| 22 | struct PyCustomKernelFunction { |
| 23 | PyCustomKernelFunction(mx::fast::CustomKernelFunction kernel, const char* tag) |
| 24 | : kernel_(std::move(kernel)), tag_(tag) {} |
| 25 | |
| 26 | std::vector<mx::array> operator()( |
| 27 | const std::vector<ScalarOrArray>& inputs_, |
| 28 | const std::vector<mx::Shape>& output_shapes, |
| 29 | const std::vector<mx::Dtype>& output_dtypes, |
| 30 | std::tuple<int, int, int> grid, |
| 31 | std::tuple<int, int, int> threadgroup, |
| 32 | const std::optional<std::vector<std::pair<std::string, nb::object>>>& |
| 33 | template_args_ = std::nullopt, |
| 34 | std::optional<float> init_value = std::nullopt, |
| 35 | bool verbose = false, |
| 36 | mx::StreamOrDevice s = {}) const { |
| 37 | std::vector<mx::array> inputs; |
| 38 | for (const auto& value : inputs_) { |
| 39 | inputs.push_back(to_array(value, std::nullopt)); |
| 40 | } |
| 41 | std::vector<std::pair<std::string, mx::fast::TemplateArg>> template_args; |
| 42 | if (template_args_) { |
| 43 | for (const auto& [name, value] : template_args_.value()) { |
| 44 | // Handle bool, int and dtype template args |
| 45 | if (nb::isinstance<bool>(value)) { |
| 46 | bool bool_val = nb::cast<bool>(value); |
| 47 | template_args.emplace_back(name, bool_val); |
| 48 | } else if (nb::isinstance<int>(value)) { |
| 49 | int int_val = nb::cast<int>(value); |
| 50 | template_args.emplace_back(name, int_val); |
| 51 | } else if (nb::isinstance<mx::Dtype>(value)) { |
| 52 | mx::Dtype dtype = nb::cast<mx::Dtype>(value); |
| 53 | template_args.emplace_back(name, dtype); |
| 54 | } else { |
| 55 | std::ostringstream msg; |
| 56 | msg << tag_ |
| 57 | << " Invalid template argument. Must be `mlx.core.Dtype`, `int` or `bool`."; |
| 58 | throw std::invalid_argument(msg.str()); |
| 59 | } |
| 60 | } |
| 61 | } |
| 62 | return kernel_( |
| 63 | inputs, |
| 64 | output_shapes, |
| 65 | output_dtypes, |
| 66 | grid, |
| 67 | threadgroup, |
| 68 | template_args, |
| 69 | init_value, |
| 70 | verbose, |
| 71 | s); |
| 72 | } |
| 73 | |
| 74 | mx::fast::CustomKernelFunction kernel_; |
| 75 | const char* tag_; |
| 76 | }; |
| 77 | |
| 78 | } // namespace |
| 79 | |