[XLA] Redesign: add ExecuteGraph to grpc service.

PiperOrigin-RevId: 195281004

diff --git a/tensorflow/compiler/xla/rpc/BUILD b/tensorflow/compiler/xla/rpc/BUILD
index 977f863..0d56a9a 100644 (file)
--- a/tensorflow/compiler/xla/rpc/BUILD
+++ b/tensorflow/compiler/xla/rpc/BUILD
@@ -55,7 +55,7 @@ tf_cc_test(
     deps = [
         ":grpc_stub",
         "//tensorflow/compiler/xla/client",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/tests:literal_test_util",
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",

diff --git a/tensorflow/compiler/xla/rpc/grpc_client_test.cc b/tensorflow/compiler/xla/rpc/grpc_client_test.cc
index b559ee4..10997c0 100644 (file)
--- a/tensorflow/compiler/xla/rpc/grpc_client_test.cc
+++ b/tensorflow/compiler/xla/rpc/grpc_client_test.cc
@@ -24,7 +24,7 @@ limitations under the License.
 #include "grpc++/security/credentials.h"
 
 #include "tensorflow/compiler/xla/client/client.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/rpc/grpc_stub.h"
 #include "tensorflow/compiler/xla/tests/literal_test_util.h"
 #include "tensorflow/core/lib/io/path.h"
@@ -84,7 +84,7 @@ TEST_F(GRPCClientTestBase, ItsAlive) {
 }
 
 TEST_F(GRPCClientTestBase, AxpyTenValues) {
-  ComputationBuilder builder(client_.get(), "axpy_10");
+  XlaBuilder builder("axpy_10");
   auto alpha = builder.ConstantR0<float>(3.1415926535);
   auto x = builder.ConstantR1<float>(
       {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0});

diff --git a/tensorflow/compiler/xla/rpc/grpc_service.cc b/tensorflow/compiler/xla/rpc/grpc_service.cc
index 0b100bd..ffb72fc 100644 (file)
--- a/tensorflow/compiler/xla/rpc/grpc_service.cc
+++ b/tensorflow/compiler/xla/rpc/grpc_service.cc
@@ -75,6 +75,13 @@ namespace xla {
       [this, arg, result]() { return service_->Execute(arg, result); });
 }
 
+::grpc::Status GRPCService::ExecuteGraph(::grpc::ServerContext* /*context*/,
+                                         const ExecuteGraphRequest* arg,
+                                         ExecuteResponse* result) {
+  return DelegateRPC(
+      [this, arg, result]() { return service_->ExecuteGraph(arg, result); });
+}
+
 ::grpc::Status GRPCService::ExecuteAsync(::grpc::ServerContext* context,
                                          const ExecuteAsyncRequest* arg,
                                          ExecuteAsyncResponse* result) {

diff --git a/tensorflow/compiler/xla/rpc/grpc_service.h b/tensorflow/compiler/xla/rpc/grpc_service.h
index fad7437..50f0279 100644 (file)
--- a/tensorflow/compiler/xla/rpc/grpc_service.h
+++ b/tensorflow/compiler/xla/rpc/grpc_service.h
@@ -54,6 +54,10 @@ class GRPCService : public grpc::XlaService::Service {
                          const ExecuteRequest* arg,
                          ExecuteResponse* result) override;
 
+  ::grpc::Status ExecuteGraph(::grpc::ServerContext* context,
+                              const ExecuteGraphRequest* arg,
+                              ExecuteResponse* result) override;
+
   ::grpc::Status ExecuteAsync(::grpc::ServerContext* context,
                               const ExecuteAsyncRequest* arg,
                               ExecuteAsyncResponse* result) override;

diff --git a/third_party/libxsmm.BUILD b/third_party/libxsmm.BUILD
index 4124f2d..78ed1f4 100644 (file)
--- a/third_party/libxsmm.BUILD
+++ b/third_party/libxsmm.BUILD
@@ -38,8 +38,8 @@ genrule(
         ":libxsmm_interface",
     ],
     visibility = [
-        "//tensorflow/core/kernels:__pkg__",
         "//third_party/eigen3:__pkg__",
+        "//tensorflow/core/kernels:__pkg__",
     ],
 )