Which item specifically is missing? mlx-swift has custom metal kernels, what else is missing? Thanks!

API like torch.autograd.Function, a way to define both forward and backward, and save some state for backward.

For the context, I am trying to implement Gaussian Splatting using mlx, and more specifically I am porting logics in https://github.com/graphdeco-inria/diff-gaussian-rasterization/blob/9c5c2028f6fbee2be239bc4c9421ff894fe4fbe0/diff_gaussian_rasterization/__init__.py#L44-L141 using mlx.

3DGS uses a PyTorch Extension written in C++ and CUDA. The forward and backward functions, radix sort are implemented in CUDA, additionally includes glm as a third parity.

Assume that you write a metal compute shader code, I suggest using the MLX C++ custom extension. You can refer to the mlx examples in mlx/examples/extension, axpby, and CMakeLists.txt. Use python install . generates build/lib.macos-.../mlx_sample_extensions/mlx_ext.metallib.

Your current options: mlx-c, mlx-swift and mlx python use fast kernel, and mlx c++ custom kernel with nanobind to python.

I have tried to write mlx c++ and c++ custom kernel code and then import it into my iOS app project. Modify this mlx-swift repo, add a new target in Package.swift. The only problem is these metal code can not use #include "mlx/backend/metal/kernels/utils.h" just like apbxy.metal.

Directory structure:

Source/
    Cmlx/
          mlx 
    MyExtension/
          Shaders/
               dummy.metal
          SwiftEntry/
                entry.h
                entry.cpp
          Dummy/
               dummy.h
               dummy.cpp

Package.swift:

.library(name: "MyMLXExtension", targets: ["MyExtension"]),
// ...
targets: [
        .target(
            name: "MyExtension",
            dependencies: ["Cmlx"],
            resources: [.process("Shaders")], // Your metal code under Source/MyExtension/Shaders directory, they will be bundled into mlx-swift_MyExtension.bundle
            publicHeadersPath: "SwiftEntry", // expose some c++ functions to Swift.
            cSettings: [
                .headerSearchPath("mlx"),
                .headerSearchPath("mlx-c"),
            ],

            cxxSettings: [
                // Copied from Cmlx target, not sure if needed.
                .headerSearchPath("mlx"),
                .headerSearchPath("mlx-c"),
                .headerSearchPath("metal-cpp"),
                .headerSearchPath("json/single_include/nlohmann"),
                .headerSearchPath("fmt/include"),
                .define("MLX_USE_ACCELERATE"),
                .define("ACCELERATE_NEW_LAPACK"),
                .define("_METAL_"),
                .define("SWIFTPM_BUNDLE", to: "\"mlx-swift_Cmlx\""),
                .define("METAL_PATH", to: "\"default.metallib\""),
                .define("MLX_VERSION", to: "\"0.24.2\""), 
            ],
            linkerSettings: [
                .linkedFramework("Foundation"),
                .linkedFramework("Metal"),
                .linkedFramework("Accelerate"),
            ]
        ),

Maybe we need some swift binding tool just like nanobind.

[Feature] custom extensions #119

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions