Closed jaro-sevcik closed 1 month ago
This is a simpler version of https://github.com/openxla/xla/pull/14807.
The main functional difference from PR#14807 is that this version also preserves argument layouts through layout canonicalization callback (even though parameter layout canonicalization seems to be disabled for GPU in OSS XLA).
[This also fixes OSS build of se_gpu_pjrt_client_test.cc
, see https://github.com/openxla/xla/pull/14899]
third_party/tensorflow/compiler/xla/pjrt/gpu:se_gpu_pjrt_client_test
is failing. Specifically the newly added MlirResultHostMemorySpaceIsSetInHloWithShardingPropagation test. Could you please fix the test?
[ RUN ] StreamExecutorGpuClientTest.MlirResultHostMemorySpaceIsSetInHloWithShardingPropagation
[third_party/tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc:1392](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc?l=1392&ws=dirkh/13446&snapshot=9): Failure
Value of: _status_or_value98.status().ok()
Actual: false
Expected: true
INVALID_ARGUMENT: No matching device found for device_id 2
=== Source Location Trace: ===
[third_party/tensorflow/compiler/xla/util.h:281](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/util.h?l=281&ws=dirkh/13446&snapshot=9)
[third_party/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h:285](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h?l=285&ws=dirkh/13446&snapshot=9)
[third_party/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:3441](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc?l=3441&ws=dirkh/13446&snapshot=9)
[third_party/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:3481](https://cs.corp.google.com/piper///depot/google3/third_party/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc?l=3481&ws=dirkh/13446&snapshot=9)
Stack trace:
0x7ffa9883f335: xla::(anonymous namespace)::StreamExecutorGpuClientTest_MlirResultHostMemorySpaceIsSetInHloWithShardingPropagation_Test::TestBody() @ ??:??
0x7ff88f738e62: testing::Test::Run() @ ??:??
0x7ff88f739ef9: testing::TestInfo::Run() @ ??:??
... Google Test internal frames ...
[ FAILED ] StreamExecutorGpuClientTest.MlirResultHostMemorySpaceIsSetInHloWithShardingPropagation (1178 ms)
In MLIR flavor of the PjRtStreamExecutorClient::Compile method, we now transfer the argument layouts and result layout from MLIR code to compile options.
If compile options already specified argument layouts, we ignore layouts from MLIR.
We also make sure that the argument/result layouts are preserved when SPMD needs to canonicalize layouts after resharding parameters and/or layouts.