0

我目前正在修改 TensorFlow 的 XLA 代码，为 GPU 编译器添加优化通道。我添加了一个std::map作为成员变量的类，用于在其他优化过程中存储和使用信息。

添加的代码如下：

#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_

该文件位于tensorflow/compiler/xla/service/gpu

错误信息如下：

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

在添加其他优化通道（例如图形重写通道）时，我没有遇到这种错误。HloInstruciton但是，我在修改类时遇到了类似的错误。

我检查了 Bazel 构建文件，我相信我已经添加了每个 cc_library 进行编译。

我正在使用 TF 2.4.1，搭配 AMD EPYC 7452 32 核处理器，GPU 以 RTX 2080 Ti * 4 为目标。

有没有办法为 TensorFlow 编译额外的实现？

4

0 回答 0