mirror of
https://github.com/ml-explore/mlx.git
synced 2025-07-23 18:11:17 +08:00
610af352d4
* Dispatch bf16 at run time when using the JIT * fix extension * fix extension build * fix extension build * Update utils.h
32 lines
721 B
Bash
32 lines
721 B
Bash
#!/bin/bash
|
|
#
|
|
# This script generates a C++ function that provides the Metal unary and binary
|
|
# ops at runtime for use with kernel generation.
|
|
#
|
|
# Copyright © 2023-24 Apple Inc.
|
|
|
|
OUTPUT_DIR=$1
|
|
CC=$2
|
|
SRC_DIR=$3
|
|
SRC_FILE=$4
|
|
CFLAGS=$5
|
|
SRC_NAME=$(basename -- "${SRC_FILE}")
|
|
JIT_INCLUDES=${SRC_DIR}/mlx/backend/metal/kernels/jit
|
|
INPUT_FILE=${SRC_DIR}/mlx/backend/metal/kernels/${SRC_FILE}.h
|
|
OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
|
|
|
|
mkdir -p "$OUTPUT_DIR"
|
|
CONTENT=$($CC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P "$INPUT_FILE" $CFLAGS 2>/dev/null)
|
|
|
|
cat << EOF > "$OUTPUT_FILE"
|
|
namespace mlx::core::metal {
|
|
|
|
const char* $SRC_NAME() {
|
|
return R"preamble(
|
|
$CONTENT
|
|
)preamble";
|
|
}
|
|
|
|
} // namespace mlx::core::metal
|
|
EOF
|