Cherry-pick from upstream to fix handling of constant resources

Summary:
This commit cherry-pick:
```
    [TF2XLA] Support must-be-constant resource variables for compilation

    Performs an explicit copy at runtime from device to host if needed.

    PiperOrigin-RevId: 341491694
    Change-Id: If4a6c0c76a1110637a06e96595c6013c8fac17e5
```
```
Remove `platform` field from shaped buffer.

    This further simplifies ShapedBuffer object as itself doesn't have any
    logic to use the platform field.

    Notice that this cl also removed some sanity check in allocation
    tracker. we can add that sanity check back if need --- just keep track
    of `platform` inside of allocation tracker as a side map.

    PiperOrigin-RevId: 339938197
    Change-Id: I090e603927ed3fccdb51254f972b3af2e1ec1470
```
```
[TF2XLA] [NFC] Provide a more informative error message when encountering a must-constant in a resource variable

    PiperOrigin-RevId: 339779174
    Change-Id: Ic364413f6794393cc15d1ccab5a5cb127560f2d1
```

With additional merge conflict fixes/API Changes fixes.

This change is required for shape inference.

Fix T43533

TF2.4 Only

Test Plan: CI

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, hakons, jakeh

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, hakons

Maniphest Tasks: T43533

Differential Revision: https://phabricator.sourcevertex.net/D49203

Author: georgepaw

tensorflow/compiler/jit/get_compiler_ir.cc

@@ -115,7 +115,7 @@ xla::StatusOr<std::string> GetCompilerIr(
 
   xla::StatusOr<std::vector<XlaCompiler::Argument>> args =
       XlaComputationLaunchContext::BuildXlaCompilerArguments(
-          constant_arg_indices, inputs, variable_infos);
+          constant_arg_indices, inputs, variable_infos, dev);
   TF_RETURN_IF_ERROR(args.status());
 
   switch (stage) {

tensorflow/compiler/jit/kernels/xla_ops.cc

@@ -213,7 +213,8 @@ static Status CompileToLocalExecutable(
 
   xla::StatusOr<std::vector<XlaCompiler::Argument>> args =
       XlaComputationLaunchContext::BuildXlaCompilerArguments(
-          constants, inputs, variable_infos, mangled_input_names);
+          constants, inputs, variable_infos,
+          static_cast<Device*>(ctx->device()), mangled_input_names);
   TF_RETURN_IF_ERROR(args.status());
   return cache->Compile(options, function, *args, compile_options,
                         lazy ? XlaCompilationCache::CompileMode::kLazy
@@ -252,8 +253,6 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
   se::Stream* stream =
       ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr;
 
-  VLOG(1) << "Executing XLA Computation...";
-
   absl::optional<se::TfAllocatorAdapter> tf_allocator_adapter;
   se::DeviceMemoryAllocator* allocator = GetAllocator(
       &tf_allocator_adapter, ctx->device(),

tensorflow/compiler/jit/xla_compilation_cache.cc

@@ -139,6 +139,7 @@ XlaCompilationCache::BuildSignature(
   for (const XlaCompiler::Argument& arg : args) {
     switch (arg.kind) {
       case XlaCompiler::Argument::kConstant:
+      case XlaCompiler::Argument::kConstantResource:
         signature.arg_values.push_back(arg.constant_value);
         break;
       case XlaCompiler::Argument::kParameter:
@@ -488,6 +489,7 @@ Status XlaCompilationCache::CompileImpl(
           argument_input_indices.push_back(i);
           break;
         }
+        case XlaCompiler::Argument::kConstantResource:
         case XlaCompiler::Argument::kResource: {
           resource_input_indices.push_back(i);
           resource_input_initialized.push_back(arg.initialized);

tensorflow/compiler/jit/xla_compile_on_demand_op.cc

@@ -152,7 +152,8 @@ Status XlaCompileOnDemandOp::Compile(
         ctx, variables_indices, variable_infos, variable_args));
 
     args = XlaComputationLaunchContext::BuildXlaCompilerArguments(
-        constant_input_indices, inputs, variable_infos);
+        constant_input_indices, inputs, variable_infos,
+        static_cast<Device*>(ctx->device()));
     TF_RETURN_IF_ERROR(args.status());
   }
 

tensorflow/compiler/jit/xla_launch_util.cc

@@ -426,7 +426,7 @@ Status XlaComputationLaunchContext::PopulateOutputs(
     ShapedBuffer buffer(
         xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_host_shape()}),
         xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_device_shape()}),
-        output.platform(), output.device_ordinal());
+        output.device_ordinal());
     buffer.buffers().CopySubtreeFrom(nontuple_buffer.buffers(),
                                      /*source_base_index=*/{},
                                      /*target_base_index=*/{0});
@@ -564,12 +564,27 @@ xla::StatusOr<std::vector<XlaCompiler::Argument>>
 XlaComputationLaunchContext::BuildXlaCompilerArguments(
     absl::Span<int const> must_be_constant_idxs,
     absl::Span<const Tensor* const> inputs,
-    absl::Span<VariableInfo const> variable_args,
+    absl::Span<VariableInfo const> variable_args, Device* device,
     const std::vector<std::string>& mangled_input_names) {
   CHECK(absl::c_is_sorted(must_be_constant_idxs));
   std::vector<XlaCompiler::Argument> out;
   out.resize(inputs.size());
 
+  // TODO(cheshire): Avoid duplication with framework/op_kernel.h
+  DeviceContext* device_context = nullptr;
+  TF_RETURN_IF_ERROR(device->TryGetDeviceContext(&device_context));
+  bool using_default_context = false;
+  auto cleanup = xla::MakeCleanup([&] {
+    if (device_context != nullptr && !using_default_context) {
+      device_context->Unref();
+    }
+  });
+  if (device_context == nullptr) {
+    using_default_context = true;
+    auto* dev_info = device->tensorflow_gpu_device_info();
+    if (dev_info) device_context = dev_info->default_context;
+  }
+
   absl::flat_hash_map<int, const VariableInfo*> variable_info_lookup;
   for (const VariableInfo& info : variable_args) {
     CHECK(!info.var() || info.lock_held())
@@ -589,14 +604,7 @@ XlaComputationLaunchContext::BuildXlaCompilerArguments(
       arg.name = mangled_input_names.at(input_num);
     }
 
-    if (absl::c_binary_search(must_be_constant_idxs, input_num)) {
-      // Handles compile-time constants.
-      TF_RET_CHECK(input->dtype() != DT_RESOURCE);
-      arg.kind = XlaCompiler::Argument::kConstant;
-      arg.type = input->dtype();
-      arg.shape = input->shape();
-      arg.constant_value = *input;
-    } else if (variable_info_lookup.count(input_num)) {
+    if (variable_info_lookup.count(input_num)) {
       // Handles resource variables.
       TF_RET_CHECK(input->dtype() == DT_RESOURCE);
       const VariableInfo& variable = *variable_info_lookup[input_num];
@@ -617,6 +625,25 @@ XlaComputationLaunchContext::BuildXlaCompilerArguments(
         arg.type = DT_INVALID;
         arg.shape = TensorShape();
       }
+
+      if (absl::c_binary_search(must_be_constant_idxs, input_num)) {
+        TF_RET_CHECK(variable.var() && variable.var()->is_initialized);
+        const Tensor* value = variable.var()->tensor();
+        Tensor value_on_host(value->dtype(), value->shape());
+        if (!device_context) {
+          value_on_host = *value;
+        } else {
+          TF_RETURN_IF_ERROR(device_context->CopyDeviceTensorToCPUSync(
+              value, "", device, &value_on_host));
+        }
+        arg.kind = XlaCompiler::Argument::kConstantResource;
+        arg.constant_value = value_on_host;
+      }
+    } else if (absl::c_binary_search(must_be_constant_idxs, input_num)) {
+      arg.kind = XlaCompiler::Argument::kConstant;
+      arg.type = input->dtype();
+      arg.shape = input->shape();
+      arg.constant_value = *input;
     } else {
       // Normal inputs.
       TF_RET_CHECK(input->dtype() != DT_RESOURCE);

tensorflow/compiler/jit/xla_launch_util.h

@@ -144,7 +144,7 @@ class XlaComputationLaunchContext {
   BuildXlaCompilerArguments(
       absl::Span<int const> must_be_constant_idxs,
       absl::Span<const Tensor* const> inputs,
-      absl::Span<VariableInfo const> variable_args,
+      absl::Span<VariableInfo const> variable_args, Device* device,
       const std::vector<std::string>& mangled_input_names = {});
 
   // Add all inputs within `ctx` as XLA arguments (returned by arguments()).

tensorflow/compiler/plugin/poplar/kernels/application_runtime/application_compile.cc

@@ -99,7 +99,8 @@ xla::StatusOr<xla::LocalExecutable*> CompileExecutable(
   TF_ASSIGN_OR_RETURN(
       std::vector<XlaCompiler::Argument> arguments,
       XlaComputationLaunchContext::BuildXlaCompilerArguments(
-          constants, inputs, variable_infos, mangled_input_names));
+          constants, inputs, variable_infos,
+          static_cast<Device*>(ctx->device()), mangled_input_names));
 
   const XlaCompiler::CompilationResult* compilation_result;
   xla::LocalExecutable* executable;

tensorflow/compiler/plugin/poplar/tests/graph_compile_io_map_test.cc

@@ -436,11 +436,11 @@ main {
 
   Shape shape = ShapeUtil::MakeShape(F32, {2});
 
-  ShapedBuffer arg0(shape, shape, platform, 0);
+  ShapedBuffer arg0(shape, shape, 0);
   arg0.set_buffer(buf0, {});
-  ShapedBuffer arg1(shape, shape, platform, 0);
+  ShapedBuffer arg1(shape, shape, 0);
   arg1.set_buffer(buf1, {});
-  ShapedBuffer arg2(shape, shape, platform, 0);
+  ShapedBuffer arg2(shape, shape, 0);
   arg2.set_buffer(buf2, {});
 
   std::vector<const ShapedBuffer*> args = {&arg0, &arg1, &arg2};
@@ -549,9 +549,9 @@ main {
 
   Shape shape = ShapeUtil::MakeShape(F32, {2});
 
-  ShapedBuffer arg0(shape, shape, platform, 0);
+  ShapedBuffer arg0(shape, shape, 0);
   arg0.set_buffer(buf, {});
-  ShapedBuffer arg1(shape, shape, platform, 0);
+  ShapedBuffer arg1(shape, shape, 0);
   arg1.set_buffer(buf, {});
 
   std::vector<const ShapedBuffer*> args = {&arg0, &arg1};

tensorflow/compiler/plugin/poplar/tests/variable_test.py

@@ -19,7 +19,11 @@
 import test_utils as tu
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.platform import googletest
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.ipu.config import IPUConfig
 from tensorflow.python.ops import array_ops
@@ -796,6 +800,74 @@ def testNonModifiedResourceIsNotOverwrittenInPlaceOp(self):
           [], "w should be copied to device once and "
           "that should be the only io event")
 
+  def testGetConstantOutOfResourceVariable(self):
+    with ops.device("/device:IPU:0"):
+
+      # Use floats to force device placement.
+      a = variables.Variable(50.0)
+      b = variables.Variable(2.0)
+
+      @def_function.function(experimental_compile=True)
+      def f(x):
+        return array_ops.reshape(
+            x,
+            [math_ops.cast(a, dtypes.int32),
+             math_ops.cast(b, dtypes.int32)])
+
+      # OK since the value is known at compile time.
+      out = f(random_ops.random_normal([10, 10]))
+      self.assertEqual(out.shape[0], 50)
+      self.assertEqual(out.shape[1], 2)
+
+  def testGetConstantOutOfResourceVariableAfterWrite(self):
+    with ops.device("/device:IPU:0"):
+
+      # Use floats to force device placement.
+      a = variables.Variable(50.0)
+      b = variables.Variable(2.0)
+
+      @def_function.function(experimental_compile=True)
+      def f(x, val1, val2):
+        a.assign(math_ops.cast(val1, dtypes.float32))
+        b.assign(math_ops.cast(val2, dtypes.float32))
+        return array_ops.reshape(
+            x,
+            [math_ops.cast(a, dtypes.int32),
+             math_ops.cast(b, dtypes.int32)])
+
+      val1 = constant_op.constant(2)
+      val2 = constant_op.constant(50)
+
+      # Returns an error, since the value known at compile time was overriden.
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  'concrete values at compile time'):
+        f(random_ops.random_normal([10, 10]), val1, val2)
+
+  def testGetConstantOutOfResourceVariableBeforeWrite(self):
+    with ops.device("/device:IPU:0"):
+
+      # Use floats to force device placement.
+      a = variables.Variable(50.0)
+      b = variables.Variable(2.0)
+
+      @def_function.function(experimental_compile=True)
+      def f(x, val1, val2):
+        out = array_ops.reshape(
+            x,
+            [math_ops.cast(a, dtypes.int32),
+             math_ops.cast(b, dtypes.int32)])
+        a.assign(math_ops.cast(val1, dtypes.float32))
+        b.assign(math_ops.cast(val2, dtypes.float32))
+        return out
+
+      val1 = constant_op.constant(2)
+      val2 = constant_op.constant(50)
+
+      # OK since the write happens after the reshape.
+      out = f(random_ops.random_normal([10, 10]), val1, val2)
+      self.assertEqual(out.shape[0], 50)
+      self.assertEqual(out.shape[1], 2)
+
 
 if __name__ == "__main__":
   os.environ['TF_XLA_FLAGS'] = ('--tf_xla_min_cluster_size=1 ' +

tensorflow/compiler/tf2xla/graph_compiler.cc

@@ -73,7 +73,7 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph,
     switch (expressions[i]->kind()) {
       case XlaExpression::Kind::kConstant:
         arg.kind = XlaCompiler::Argument::kConstant;
-        arg.constant_value = expressions[i]->constant_value();
+        arg.constant_value = *expressions[i]->constant_value();
         break;
       case XlaExpression::Kind::kXlaOp:
         if (arg_must_be_compile_time_constant[i]) {