zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-18 07:18:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add dispatchThreads to custom kernel doc (#1551)

Browse Source 
* add dispatchThreads info

* update

* add link
This commit is contained in:
Alex Barron
2024-11-01 13:07:48 -07:00
committed by GitHub
parent eac961ddb1
commit 9e516b71ea
2 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/src/fast.cpp
View File

@@ -278,7 +278,9 @@ void init_fast(nb::module_& parent_module) {
output_shapes (List[Sequence[int]]): The list of shapes for each output in ``output_names``.
output_dtypes (List[Dtype]): The list of data types for each output in ``output_names``.
grid (tuple[int, int, int]): 3-tuple specifying the grid to launch the kernel with.
This will be passed to ``MTLComputeCommandEncoder::dispatchThreads``.
threadgroup (tuple[int, int, int]): 3-tuple specifying the threadgroup size to use.
This will be passed to ``MTLComputeCommandEncoder::dispatchThreads``.
template (List[Tuple[str, Union[bool, int, Dtype]]], optional): Template arguments.
These will be added as template arguments to the kernel definition. Default: ``None``.
init_value (float, optional): Optional value to use to initialize all of the output arrays.
@@ -300,6 +302,8 @@ void init_fast(nb::module_& parent_module) {
R"pbdoc(
A jit-compiled custom Metal kernel defined from a source string.
Full documentation: :ref:`custom_metal_kernels`.
Args:
name (str): Name for the kernel.
input_names (List[str]): The parameter names of the inputs in the