Add support for parallel quantization on Mac

This commit is contained in:
duzx16 2023-05-04 21:45:31 +02:00
parent 63d66b0572
commit f6b88da8c1
1 changed files with 7 additions and 2 deletions

View File

@ -7,6 +7,7 @@ import bz2
import torch
import base64
import ctypes
import sys
from transformers.utils import logging
from typing import List
@ -142,8 +143,12 @@ class CPUKernel:
kernel_file = source_code[:-2] + ".so"
if compile_parallel_kernel:
if sys.platform != 'darwin':
compile_command = "gcc -O3 -fPIC -pthread -fopenmp -std=c99 {} -shared -o {}".format(
source_code, kernel_file)
else:
compile_command = "clang -O3 -fPIC -pthread -Xclang -fopenmp -lomp -std=c99 {} -shared -o {}".format(
source_code, kernel_file)
print("Compiling", compile_command)
exit_state = os.system(compile_command)
if not exit_state: