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

[CUDA] Bundle CCCL for JIT compilation (#2357)

Browse Source 
* Ship CCCL for JIT compilation

* Remove cexpf
This commit is contained in:
Cheng
2025-07-12 10:45:37 +09:00
committed by GitHub
parent 42cc9cfbc7
commit 6325f60d52
9 changed files with 48 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
mlx/backend/metal/device.cpp
View File

@@ -1,20 +1,18 @@
// Copyright © 2023-2024 Apple Inc.
#include <cstdlib>
#include <filesystem>
#include <sstream>
#define NS_PRIVATE_IMPLEMENTATION
#define CA_PRIVATE_IMPLEMENTATION
#define MTL_PRIVATE_IMPLEMENTATION
#include "mlx/backend/common/utils.h"
#include "mlx/backend/metal/device.h"
#include "mlx/backend/metal/metal.h"
#include "mlx/backend/metal/utils.h"
#include "mlx/utils.h"
namespace fs = std::filesystem;
namespace mlx::core::metal {
namespace {
@@ -80,12 +78,7 @@ MTL::Library* try_load_bundle(
std::pair<MTL::Library*, NS::Error*> load_colocated_library(
MTL::Device* device,
const std::string& relative_path) {
std::string binary_dir = get_binary_directory();
if (binary_dir.size() == 0) {
return {nullptr, nullptr};
}
auto path = fs::path(binary_dir) / relative_path;
auto path = current_binary_dir() / relative_path;
if (!path.has_extension()) {
path.replace_extension(".metallib");
}
@@ -197,7 +190,7 @@ MTL::Library* load_library(
std::ostringstream msg;
msg << "Failed to load the metallib " << lib_name << ".metallib. "
<< "We attempted to load it from <" << get_binary_directory() << "/"
<< "We attempted to load it from <" << current_binary_dir() << "/"
<< lib_name << ".metallib" << ">";
#ifdef SWIFTPM_BUNDLE
msg << " and from the Swift PM bundle.";

16
mlx/backend/metal/device.h
View File

@@ -3,8 +3,6 @@
#pragma once
#include <Metal/Metal.hpp>
#include <dlfcn.h>
#include <filesystem>
#include <functional>
#include <mutex>
#include <shared_mutex>
@@ -15,22 +13,8 @@
#include "mlx/array.h"
#include "mlx/device.h"
namespace fs = std::filesystem;
namespace mlx::core::metal {
// Note, this function must be left inline in a header so that it is not
// dynamically linked.
inline std::string get_binary_directory() {
Dl_info info;
std::string directory;
int success = dladdr((void*)get_binary_directory, &info);
if (success) {
directory = fs::path(info.dli_fname).remove_filename().c_str();
}
return directory;
}
using MTLFCList =
std::vector<std::tuple<const void*, MTL::DataType, NS::UInteger>>;