zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 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/common/utils.cpp
View File

@@ -1,5 +1,7 @@
// Copyright © 2023-2024 Apple Inc.
#include <dlfcn.h>
#include "mlx/backend/common/utils.h"
#include "mlx/primitives.h"
@@ -11,6 +13,17 @@ std::string get_primitive_string(Primitive* primitive) {
return op_t.str();
}
std::filesystem::path current_binary_dir() {
static std::filesystem::path binary_dir = []() {
Dl_info info;
if (!dladdr(reinterpret_cast<void*>(&current_binary_dir), &info)) {
throw std::runtime_error("Unable to get current binary dir.");
}
return std::filesystem::path(info.dli_fname).parent_path();
}();
return binary_dir;
}
std::tuple<Shape, std::vector<Strides>> collapse_contiguous_dims(
const Shape& shape,
const std::vector<Strides>& strides,

4
mlx/backend/common/utils.h
View File

@@ -2,6 +2,7 @@
#pragma once
#include <filesystem>
#include <tuple>
#include <vector>
@@ -11,6 +12,9 @@ namespace mlx::core {
std::string get_primitive_string(Primitive* primitive);
// Return the directory that contains current shared library.
std::filesystem::path current_binary_dir();
inline int64_t
elem_to_loc(int elem, const Shape& shape, const Strides& strides) {
int64_t loc = 0;