What is need to be done when adding additional code in TensorFlow?

I am currently working with TensorFlow XLA to add a couple of optimization passes that rewrite some HLO nodes to CustomCall.
The implemented optimization pass along works fine, but when I added some headers, which contain classes used in the optimization pass, building TensorFlow results with undefined symbol error.

The added code is as the following:

#ifndef _TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NDPX_DEVICE_TABLE_H_
#define _TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NDPX_DEVICE_TABLE_H_

#include <map>

namespace xla {
namespace gpu {

enum class Device {
  GPU,
  PIM,
  UNDEFINED
};

class DeviceAssignTable {
 private:
  std::map<int, Device> device_table_;
  
  DeviceAssignTable() {}
  static DeviceAssignTable* deviceAssignTable;

 public:
  void AssignDevice(int hlo_id, Device device) {
    device_table_[hlo_id] = device;
  }

  Device FindAssignInfo(int hlo_id) {
    return device_table_[hlo_id];
  }

  void operator=(const DeviceAssignTable&) = delete;

  static DeviceAssignTable* GetInstance() {
    if (deviceAssignTable == nullptr) {
      deviceAssignTable = new DeviceAssignTable();
    }
    return deviceAssignTable;
  }
};

static DeviceAssignTable* deviceAssignTable = DeviceAssignTable::GetInstance();

} // namespace gpu
} // namespace xla

#endif  // _TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NDPX_DEVICE_TABLE_H_

I am planning to use this class in one of my optimization passes I made, yet it ends in undefined symbol error.
I added all the files I have created to BUILD file in tensorflow/compiler/xla/service/gpu (where the file exists).
Is adding a global variable not permitted in TensorFlow? Is there something additional to be done when adding a class like this?

Thanks for spending your time to this question.

System information

  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04): CentOS 7
  • Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
  • TensorFlow installed from (source or binary): source
  • TensorFlow version: 2.4.1
  • Python version: 3.6.13
  • Installed using virtualenv? pip? conda?: conda
  • Bazel version (if compiling from source): 3.7.1
  • GCC/Compiler version (if compiling from source): 8.3.0
  • CUDA/cuDNN version: 10.1 / 7.6.5
  • GPU model and memory: 2080Ti * 4, 11019Mib per GPU

Provide the exact sequence of commands / steps that you executed before running into the problem

Build command:

bazel build --config=cuda //tensorflow/tools/pip_package:build_pip_package

Any other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

ERROR: /home/jueonpark/cxl-simulator/tensorflow/tensorflow/python/keras/api/BUILD:111:19: Executing genrule //tensorflow/python/keras/api:keras_python_api_gen failed (Exit 1): bash failed: error executing command /bin/bash -c ... (remaining 1 argument(s) skipped)
Traceback (most recent call last):
  File "/home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/pywrap_tensorflow.py", line 64, in <module>
    from tensorflow.python._pywrap_tensorflow_internal import *
ImportError: /home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/_pywrap_tensorflow_internal.so: undefined symbol: _ZN3xla3gpu17DeviceAssignTable17deviceAssignTableE

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/tools/api/generator/create_python_api.py", line 26, in <module>
    from tensorflow.python.tools.api.generator import doc_srcs
  File "/home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/__init__.py", line 39, in <module>
    from tensorflow.python import pywrap_tensorflow as _pywrap_tensorflow
  File "/home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/pywrap_tensorflow.py", line 83, in <module>
    raise ImportError(msg)
ImportError: Traceback (most recent call last):
  File "/home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/pywrap_tensorflow.py", line 64, in <module>
    from tensorflow.python._pywrap_tensorflow_internal import *
ImportError: /home/jueonpark/.cache/bazel/_bazel_jueonpark/a1d79b17c57f21a0bc8c6679dc05b873/execroot/org_tensorflow/bazel-out/host/bin/tensorflow/python/keras/api/create_tensorflow.python_api_keras_python_api_gen.runfiles/org_tensorflow/tensorflow/python/_pywrap_tensorflow_internal.so: undefined symbol: _ZN3xla3gpu17DeviceAssignTable17deviceAssignTableE


Failed to load the native TensorFlow runtime.

See https://www.tensorflow.org/install/errors

for some common reasons and solutions.  Include the entire stack trace
above this error message when asking for help.
Target //tensorflow/tools/pip_package:build_pip_package failed to build
Use --verbose_failures to see the command lines of failed build steps.
ERROR: /home/jueonpark/cxl-simulator/tensorflow/tensorflow/tools/pip_package/BUILD:69:10 Executing genrule //tensorflow/python/keras/api:keras_python_api_gen failed (Exit 1): bash failed: error executing command /bin/bash -c ... (remaining 1 argument(s) skipped)
INFO: Elapsed time: 187.897s, Critical Path: 120.74s
INFO: 2541 processes: 120 internal, 2421 local.
FAILED: Build did NOT complete successfully

Sorry for taking your time. The problem was linking error, and it was solved by defining functions in *.cc file.