experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Provide a way to test GPU buffer assignment, test buffer assignment for GEMM rewrite

Browse Source 
PiperOrigin-RevId: 254607491
This commit is contained in:
George Karpenkov 2019-06-22 23:42:00 -07:00 committed by TensorFlower Gardener
parent d4b6a19ee7
commit eeda570f69
2 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/xla/service/gpu/gpu_executable.h
View File

@ -89,6 +89,10 @@ class GpuExecutable : public Executable {
const ServiceExecutableRunOptions* run_options,
absl::Span<const ShapedBuffer* const> arguments) override;
std::shared_ptr<const BufferAssignment> GetBufferAssignment() const {
return assignment_;
}
private:
StatusOr<ScopedShapedBuffer> Execute(
const ServiceExecutableRunOptions* run_options,

1
tensorflow/compiler/xla/service/gpu/tests/BUILD
View File

@ -58,6 +58,7 @@ tf_cc_test(
deps = [
":gpu_codegen_test",
"//tensorflow/compiler/xla:debug_options_flags",
"//tensorflow/compiler/xla:statusor",
"//tensorflow/compiler/xla/service:gpu_plugin",
"//tensorflow/compiler/xla/service:hlo",
"//tensorflow/compiler/xla/service:hlo_module_config",