From d0f97926135a9c85ff9a32189644124a60c5ac99 Mon Sep 17 00:00:00 2001
From: Derui Yang <ydrml@hotmail.com>
Date: Mon, 16 Oct 2023 19:53:28 +0800
Subject: [PATCH] Fix: add building option for NNet (#162)

Signed-off-by: YdrMaster <ydrml@hotmail.com>
---
 CMakeLists.txt              | 15 +++++++++++----
 src/kernels/cpu/membound.cc |  4 ++++
 src/operators/membound.cc   |  4 ++++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 72df016c..d942fcd4 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -5,6 +5,7 @@ option(USE_KUNLUN "Support KUNLUN XPU" OFF)
 option(USE_INTELCPU "Support INTELCPU" OFF)
 option(USE_BACKTRACE "Print backtrace on exception and segmentation fault" ON)
 option(USE_PROTOBUF "Serialize and deserialize tensors" OFF)
+option(BUILD_NNET "Build nnet" OFF)
 option(BUILD_DIST "Build project for distributed running" OFF)
 option(BUILD_TEST "Build tests" OFF)
 
@@ -30,14 +31,14 @@ if(CMAKE_BUILD_TYPE STREQUAL "Debug")
     add_compile_definitions(DEBUG_MODE)
 elseif(CMAKE_BUILD_TYPE STREQUAL "Release")
     message("Configuring for Release build.")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O2") 
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O2")
     add_compile_definitions(NDEBUG)
 elseif(CMAKE_BUILD_TYPE STREQUAL "RelWithDebInfo")
     message("Configuring for RelWithDebInfo build.")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O2") 
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O2")
 else()
     message("Build type not specified. Configuring for RelWithDebInfo build.")
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O2") 
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O2")
 endif()
 
 
@@ -123,7 +124,13 @@ set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELWITHDEBINFO} -UNDEBUG")
 
 
 # Source files
-file(GLOB_RECURSE SRC src/ffi/*.cc src/core/*.cc src/kernels/cpu/*.cc src/nnet/*.cc src/operators/*.cc src/utils/*.cc)
+file(GLOB_RECURSE SRC src/ffi/*.cc src/core/*.cc src/kernels/cpu/*.cc src/operators/*.cc src/utils/*.cc)
+
+if(BUILD_NNET)
+  add_compile_definitions(BUILD_NNET=1)
+  file(GLOB_RECURSE SRC_NNET src/nnet/*.cc)
+  list (APPEND SRC ${SRC_NNET})
+endif()
 
 if(USE_CUDA)
   file(GLOB_RECURSE SRC_CUDA src/cuda/*.cc src/cuda/*.cu src/kernels/cuda/*.cc src/kernels/cuda/*.cu)
diff --git a/src/kernels/cpu/membound.cc b/src/kernels/cpu/membound.cc
index 18fb304c..b6b6c7ee 100644
--- a/src/kernels/cpu/membound.cc
+++ b/src/kernels/cpu/membound.cc
@@ -1,3 +1,5 @@
+#ifdef BUILD_NNET
+
 #include "operators/membound.h"
 #include "core/kernel.h"
 #include "nnet/Visitor/Interpreter.h"
@@ -82,3 +84,5 @@ REGISTER_KERNEL(Device::CPU, OpType::MemBound, DataType::UInt32,
                 MemboundInterpreter, "MemboundInterpreter_CPU");
 
 } // namespace infini
+
+#endif
diff --git a/src/operators/membound.cc b/src/operators/membound.cc
index 8fd6d294..be757b36 100644
--- a/src/operators/membound.cc
+++ b/src/operators/membound.cc
@@ -1,3 +1,5 @@
+#ifdef BUILD_NNET
+
 #include "operators/membound.h"
 #include "nnet/Visitor/CheckOOBVisitor.h"
 #include "nnet/Visitor/HashVisitor.h"
@@ -84,3 +86,5 @@ bool MemBoundObj::checkOOB(nnet::Expr expr) {
 }
 
 } // namespace infini
+
+#endif