zxdu20 commited on
Commit
f6b88da
1 Parent(s): 63d66b0

Add support for parallel quantization on Mac

Browse files
Files changed (1) hide show
  1. quantization.py +7 -2
quantization.py CHANGED
@@ -7,6 +7,7 @@ import bz2
7
  import torch
8
  import base64
9
  import ctypes
 
10
  from transformers.utils import logging
11
 
12
  from typing import List
@@ -142,8 +143,12 @@ class CPUKernel:
142
  kernel_file = source_code[:-2] + ".so"
143
 
144
  if compile_parallel_kernel:
145
- compile_command = "gcc -O3 -fPIC -pthread -fopenmp -std=c99 {} -shared -o {}".format(
146
- source_code, kernel_file)
 
 
 
 
147
  print("Compiling", compile_command)
148
  exit_state = os.system(compile_command)
149
  if not exit_state:
 
7
  import torch
8
  import base64
9
  import ctypes
10
+ import sys
11
  from transformers.utils import logging
12
 
13
  from typing import List
 
143
  kernel_file = source_code[:-2] + ".so"
144
 
145
  if compile_parallel_kernel:
146
+ if sys.platform != 'darwin':
147
+ compile_command = "gcc -O3 -fPIC -pthread -fopenmp -std=c99 {} -shared -o {}".format(
148
+ source_code, kernel_file)
149
+ else:
150
+ compile_command = "clang -O3 -fPIC -pthread -Xclang -fopenmp -lomp -std=c99 {} -shared -o {}".format(
151
+ source_code, kernel_file)
152
  print("Compiling", compile_command)
153
  exit_state = os.system(compile_command)
154
  if not exit_state: