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?
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.