mirror of
https://github.com/ml-explore/mlx.git
synced 2025-06-26 02:33:21 +08:00
42 lines
782 B
Bash
42 lines
782 B
Bash
#!/bin/bash
|
|
#
|
|
# This script generates a C++ function that provides the CPU
|
|
# code for use with kernel generation.
|
|
#
|
|
# Copyright © 2023-24 Apple Inc.
|
|
|
|
|
|
OUTPUT_FILE=$1
|
|
GCC=$2
|
|
SRCDIR=$3
|
|
CLANG=$4
|
|
ARCH=$5
|
|
|
|
if [ "$CLANG" = "TRUE" ]; then
|
|
read -r -d '' INCLUDES <<- EOM
|
|
#include <cmath>
|
|
#include <complex>
|
|
#include <cstdint>
|
|
#include <vector>
|
|
#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC
|
|
#include <arm_fp16.h>
|
|
#endif
|
|
EOM
|
|
CC_FLAGS="-arch ${ARCH} -nobuiltininc -nostdinc"
|
|
else
|
|
CC_FLAGS="-std=c++17"
|
|
fi
|
|
|
|
CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
|
|
|
|
cat << EOF > "$OUTPUT_FILE"
|
|
const char* get_kernel_preamble() {
|
|
return R"preamble(
|
|
$INCLUDES
|
|
$CONTENT
|
|
using namespace mlx::core;
|
|
using namespace mlx::core::detail;
|
|
)preamble";
|
|
}
|
|
EOF
|