Skip to content

Commit 97f11a8

Browse files
author
Avik Pal
committed
Add option to disable NNPACK
1 parent 6b50e92 commit 97f11a8

File tree

1 file changed

+16
-9
lines changed

1 file changed

+16
-9
lines changed

src/nnpack/NNPACK.jl

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,23 @@ const shared_threadpool = Ref(C_NULL)
1313

1414
@init begin
1515
check_deps()
16-
status = nnp_initialize()
17-
if status == nnp_status_unsupported_hardware
18-
@warn "HARDWARE is unsupported by NNPACK so falling back to default NNlib"
19-
else
20-
include(nnlib_interface_path)
21-
end
2216
try
23-
global NNPACK_CPU_THREADS = parse(UInt64, ENV["NNPACK_CPU_THREADS"])
17+
global ENABLE_NNPACK = parse(UInt64, ENV["ENABLE_NNPACK"])
2418
catch
25-
global NNPACK_CPU_THREADS = 4
19+
global ENABLE_NNPACK = 1
20+
end
21+
if ENABLE_NNPACK == 1
22+
status = nnp_initialize()
23+
if status == nnp_status_unsupported_hardware
24+
@warn "HARDWARE is unsupported by NNPACK so falling back to default NNlib"
25+
else
26+
include(nnlib_interface_path)
27+
end
28+
try
29+
global NNPACK_CPU_THREADS = parse(UInt64, ENV["NNPACK_CPU_THREADS"])
30+
catch
31+
global NNPACK_CPU_THREADS = 4
32+
end
33+
shared_threadpool[] = pthreadpool_create(NNPACK_CPU_THREADS)
2634
end
27-
shared_threadpool[] = pthreadpool_create(NNPACK_CPU_THREADS)
2835
end

0 commit comments

Comments
 (0)