Modifying TensorFlow code, especially XLA, yet building TF emits the following error

Issue

I am currently modifying TensorFlow’s XLA code, adding optimization pass for GPU compiler.
I have added a class with std::map as a member variable, to store and use information among other optimization passes.

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,
  NDPX
};

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

 public:
  void AssignDevice(int hlo_id, Device device);
  Device FindAssignInfo(int hlo_id);

  void operator=(const DeviceAssignTable&) = delete;

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

} // namespace gpu
} // namespace xla

#endif  // _TENSORFLOW_COMPILER_XLA_SERVICE_GPU_NDPX_DEVICE_TABLE_H_

The file is located in tensorflow/compiler/xla/service/gpu

The error message is as the following:

ERROR: /home/jueonpark/cxl-simulator/tensorflow/tensorflow/python/keras/api/BUILD:124:19: Executing genrule //tensorflow/python/keras/api:keras_python_api_gen_compat_v1 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_compat_v1.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_compat_v1.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_compat_v1.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_compat_v1.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_compat_v1.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_compat_v1.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_compat_v1.runfiles/org_tensorflow/tensorflow/python/_pywrap_tensorflow_internal.so: undefined symbol: _ZN3xla3gpu17DeviceAssignTable17deviceAssignTableE

While adding other optimizaiton pass such as graph rewriting pass, I did not face this kind of error. However, I have faced similar error while modifying HloInstruciton class.

I checked the Bazel build file, and I believe I have added every cc_library for compilation.

I am using TF 2.4.1, with AMD EPYC 7452 32-Core Processor with GPU targeting RTX 2080 Ti * 4.

Is there any way to compile additional implementations for TensorFlow?

Solution

The problem was simple linking error. To avoid the problem, just define the function on *.cc file and declare the function on corresponding header file.

Answered By – Jueon Park

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published