Use Poplar host function.

jakeh-gc · georgepaw · commit 353c4f3b8bf6 · 2021-12-02T13:45:27.000Z
Summary: What's changed: - Use Poplar Host Functions to tell poplar when we have a stream copy with a host data-dependency. - Allows `opt_flags.set("streamCallbacks.maxLookahead", "unlimited");` to be safely set by default, improving throughput. - There's a slight change to the UserOp API. - In testing, it remains ABI compatible with any old ones, if any even exist. Resolves T50668 Test Plan: CI tests - particularly user ops and host embedding tests. Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, babakk Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, babakk Subscribers: jorgec, babakk Maniphest Tasks: T50668 Differential Revision: https://phabricator.sourcevertex.net/D56667
diff --git a/tensorflow/compiler/plugin/poplar/docs/custom_codelet.rst b/tensorflow/compiler/plugin/poplar/docs/custom_codelet.rst
@@ -617,9 +617,9 @@ The signature of the callback function is:
 
   extern "C"
   void Callback(
-      const std::vector<void*>& data,
+      const std::vector<const void*>& data,
       const std::vector<std::uint32_t>& number_of_elements,
-      std::vector<void*>& outputs,
+      const std::vector<void*>& outputs,
       const std::string& attributes,
       const std::string& name);
 
diff --git a/tensorflow/compiler/plugin/poplar/driver/compiler_annotations.h b/tensorflow/compiler/plugin/poplar/driver/compiler_annotations.h
@@ -18,6 +18,8 @@ limitations under the License.
 
 #include <map>
 #include <set>
+#include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -198,28 +200,18 @@ struct StreamCopyInfo {
 // assosiated with that instruction.
 using StreamInfos = std::unordered_map<std::string, std::list<StreamCopyInfo>>;
 
-// Stream meta info contains the information relating to the setup of the output
-// streams. We need to know how many outputs there are and how much data to
-// allocate in each buffer.
-struct StreamCopyMetaInfo {
-  StreamCopyMetaInfo() {}
-  StreamCopyMetaInfo(const HloInstruction* inst, std::uint32_t input_count)
-      : parent_instruction(inst), num_inputs(input_count) {}
+struct HostFunctionInfo {
+  using FunctionType = std::function<void(const std::vector<const void*>& input,
+                                          const std::vector<void*>& outputs)>;
 
-  // The instruction the user op came from. We use this as a unique identifier
-  // for the inputs/outputs so we can sort the input/outputs by operation.
   const HloInstruction* parent_instruction;
-
-  // Track all of the output streams, we do this so we can allocate them in
-  // advance.
-  std::list<StreamCopyInfo*> output_stream_info;
-
-  // The number of inputs this operation has.
-  std::uint32_t num_inputs;
+  std::string handle;
+  std::vector<Shape> input_shapes;
+  std::vector<Shape> output_shapes;
+  FunctionType function;
 };
 
-// We track one metainfo struct for each stream copy which the user has added.
-using StreamMetaInfos = std::unordered_map<std::string, StreamCopyMetaInfo>;
+using HostFunctionInfos = std::unordered_map<std::string, HostFunctionInfo>;
 
 // This structure contains all information which we generate that pertains
 // to the XLA graph, as opposed to the poplar lowering of that graph.
@@ -236,8 +228,6 @@ struct CompilerAnnotations {
 
   StreamInfos stream_infos;
 
-  StreamMetaInfos stream_meta_infos;
-
   SendRecvInfos send_infos;
   SendRecvInfos recv_infos;
 
@@ -265,6 +255,9 @@ struct CompilerAnnotations {
   OutputInfos entry_output_infos;
   // Feed output descriptions.
   OutputInfos feed_output_infos;
+
+  // Host function information.
+  HostFunctionInfos host_function_infos;
 };
 
 inline Status AddInfeedInfo(CompilerAnnotations& compiler_annotations,
diff --git a/tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/host_embedding.cc b/tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/host_embedding.cc
@@ -130,25 +130,13 @@ class HostEmbeddingLookupOp : public PoplarOpDef {
         {inst->name(), inst->EmbeddingId(), inst->operand(0)->shape(),
          output_shape, inst->SplittingStrategy()});
 
-    auto index_buffer = graph.addDeviceToHostFIFO(
-        inst->name() + inst->EmbeddingId() + "_indices", indices.elementType(),
-        indices.numElements());
+    auto lookup_fn = graph.addHostFunction(
+        /* handle = */ inst->name() + inst->EmbeddingId(),
+        /* inputs = */ {{indices.elementType(), indices.numElements()}},
+        /* outputs = */ {{output.elementType(), output.numElements()}});
 
-    auto activation_fifo = graph.addHostToDeviceFIFO(
-        inst->name() + inst->EmbeddingId() + "_activations",
-        output.elementType(), output.numElements());
-
-    // Send the indices to the host.
-    seq.add(poplar::program::Copy(indices, index_buffer, false,
-                                  {debug_name_and_id}));
-
-    // Sync to avoid any stream merging due to host-side data dependecy.
-    seq.add(
-        poplar::program::Sync(poplar::SyncType::INTERNAL, {debug_name_and_id}));
-
-    // Read the values from the host.
-    seq.add(poplar::program::Copy(activation_fifo, output, false,
-                                  {debug_name_and_id}));
+    seq.add(poplar::program::Call(lookup_fn, {indices}, {output},
+                                  {debug_name_and_id, "call"}));
 
     TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, output));
 
@@ -417,18 +405,15 @@ class HostEmbeddingUpdateOp : public PoplarOpDef {
         {inst->name(), inst->EmbeddingId(), inst->operand(2)->shape(),
          inst->operand(1)->shape()});
 
-    auto index_buffer = graph.addDeviceToHostFIFO(
-        inst->name() + inst->EmbeddingId() + "_indices", indices.elementType(),
-        indices.numElements());
-
-    auto grad_fifo =
-        graph.addDeviceToHostFIFO(inst->name() + inst->EmbeddingId() + "_grads",
-                                  grads.elementType(), grads.numElements());
+    auto update_fn = graph.addHostFunction(
+        /* handle = */ inst->name() + inst->EmbeddingId(),
+        /* inputs = */
+        {{indices.elementType(), indices.numElements()},
+         {grads.elementType(), grads.numElements()}},
+        /* outputs = */ {});
 
-    seq.add(poplar::program::Copy(indices, index_buffer, false,
-                                  {debug_name_and_id}));
-    seq.add(
-        poplar::program::Copy(grads, grad_fifo, false, {debug_name_and_id}));
+    seq.add(poplar::program::Call(update_fn, {indices, grads}, {},
+                                  {debug_name_and_id, "call"}));
 
     return seq;
   }
diff --git a/tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/poputil/user_op.cc b/tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/poputil/user_op.cc
@@ -71,9 +71,9 @@ class UserOpImpl : public PoplarOpDef {
     // the operation on the host by reading the raw bytes, executing the on the
     // host, then copying back.
     void (*as_function_host_rw_ptr)(
-        const std::vector<void*>& data,
+        const std::vector<const void*>& data,
         const std::vector<std::uint32_t>& number_of_elements,
-        std::vector<void*>& outputs, const std::string& attributes,
+        const std::vector<void*>& outputs, const std::string& attributes,
         const std::string& debugPrefix);
 
     // Convert the function pointer to each of the types of function we could
@@ -113,61 +113,33 @@ class UserOpImpl : public PoplarOpDef {
     // stream we create the StreamCopyInfo structures to communicate with
     // poplar_executor which does the actual linking of the streams.
     if (is_user_read_write) {
-      // A wrapper around the user functor which we finally call down to.
-      auto functor = [=](std::vector<void*>& data,
-                         std::vector<std::uint32_t>& number_of_elements,
-                         std::vector<void*>& outputs) {
-        as_function_host_rw_ptr(data, number_of_elements, outputs, attributes,
-                                instruction_name);
-      };
-
-      // We add a map of user ops to their owned streams.
-      res.annotations.stream_infos.insert({instruction_name, {}});
-      std::list<StreamCopyInfo>& info_list =
-          res.annotations.stream_infos[instruction_name];
-
-      // Allocate a stream info
-      res.annotations.stream_meta_infos[instruction_name] = {inst,
-                                                             number_of_inputs};
-      StreamCopyMetaInfo& meta_info =
-          res.annotations.stream_meta_infos[instruction_name];
+      auto& host_function_info =
+          res.annotations.host_function_infos[instruction_name];
+      host_function_info.parent_instruction = user_op_inst;
+      host_function_info.handle = instruction_name;
+
+      std::vector<poplar::Tensor> inputs;
+      std::vector<poplar::Graph::HostFunctionArgument> in_args;
+      std::vector<poplar::Graph::HostFunctionArgument> out_args;
+      std::vector<std::uint32_t> in_args_elems;
+      inputs.resize(user_op_inst->NumInputs());
+      outputs.resize(number_of_outputs);
 
+      // Collect the input tensors and input arg descriptions.
       for (std::uint32_t i = 0; i < user_op_inst->NumInputs(); ++i) {
         // Get the poplar tensor.
         TF_ASSIGN_OR_RETURN(poplar::Tensor in,
                             FindInstructionInput(tensor_map, res, inst, i, seq,
                                                  {debug_info}, false));
 
-        // Give each input a stream identifier based on the instruction name.
-        const std::string stream_name =
-            instruction_name + "_read_" + std::to_string(i);
-
-        // Create a datastream for the input tensor.
-        poplar::DataStream stream = graph.addDeviceToHostFIFO(
-            stream_name, in.elementType(), in.numElements());
-
-        // Allocate this structure to communicate to the executor, which
-        // callbacks to register to which input tensors.
-        const uint32_t num_elements = static_cast<uint32_t>(in.numElements());
-        const uint32_t type_size = static_cast<uint32_t>(
-            graph.getTarget().getTypeSize(in.elementType()));
-        StreamCopyInfo info{inst,      stream_name, num_elements,
-                            type_size, i,           functor};
-        info_list.push_back(info);
-
-        // Copy from the tensor into the host stream. We will later attach a
-        // callback to this.
-        seq.add(poplar::program::Copy(in, stream, false, {debug_info}));
-      }
-
-      // Add an ontile sync to stop the copies from host being merged with the
-      // above as there is an invisble dependency in the callback.
-      seq.add(poplar::program::Sync(poplar::SyncType::INTERNAL, {debug_info}));
+        in_args.emplace_back(in.elementType(), in.numElements());
+        in_args_elems.push_back(in.numElements());
+        host_function_info.input_shapes.push_back(inst->operand(i)->shape());
 
-      outputs.resize(number_of_outputs);
+        inputs[i] = in;
+      }
 
-      // Now go over and add a copy from the device back to the host for each
-      // output.
+      // Collect the output tensors and output arg descriptions.
       for (std::uint32_t output_index = 0; output_index != number_of_outputs;
            output_index++) {
         xla::Shape shape = output_shape.tuple_shapes()[output_index];
@@ -177,36 +149,27 @@ class UserOpImpl : public PoplarOpDef {
             AddTensor(graph, TensorLocation{inst, output_index}, shape, res,
                       tensor_map, {debug_info, "output"}));
 
-        // Add stream ID for each output tensor.
-        const std::string stream_name =
-            instruction_name + "_write_" + std::to_string(output_index);
-
-        // Copy from the host into these new tensors.
-        poplar::DataStream stream =
-            graph.addHostToDeviceFIFO(stream_name, output_tensor.elementType(),
-                                      output_tensor.numElements());
-
-        // Allocate this structure to communicate to the executor so the
-        // executor knows how much memory to allocate for the callback to write
-        // into.
-        const uint32_t num_elements =
-            static_cast<uint32_t>(output_tensor.numElements());
-        const uint32_t type_size = static_cast<uint32_t>(
-            graph.getTarget().getTypeSize(output_tensor.elementType()));
-        StreamCopyInfo info{inst, stream_name, num_elements, type_size,
-                            output_index};
-        info_list.push_back(std::move(info));
-
-        // Store a reference to this stream copy info.
-        StreamCopyInfo* ref = &info_list.back();
-        meta_info.output_stream_info.push_back(ref);
-
-        // Add the copy to the graph.
-        seq.add(
-            poplar::program::Copy(stream, output_tensor, false, {debug_info}));
-
+        out_args.emplace_back(output_tensor.elementType(),
+                              output_tensor.numElements());
+        host_function_info.output_shapes.push_back(shape);
         outputs[output_index] = output_tensor;
       }
+
+      host_function_info.function = [as_function_host_rw_ptr, in_args_elems,
+                                     attributes, instruction_name](
+                                        const std::vector<const void*>& input,
+                                        const std::vector<void*>& outputs) {
+        as_function_host_rw_ptr(input, in_args_elems, outputs, attributes,
+                                instruction_name);
+      };
+
+      // Create the host function
+      auto user_fn_device =
+          graph.addHostFunction(instruction_name, in_args, out_args);
+
+      // Add the device call to the host function.
+      seq.add(
+          poplar::program::Call(user_fn_device, inputs, outputs, debug_info));
     } else {
       if (!is_gradient) {
         std::vector<poplar::Tensor> inputs(user_op_inst->NumInputs());
diff --git a/tensorflow/compiler/plugin/poplar/driver/poplar_compiler.cc b/tensorflow/compiler/plugin/poplar/driver/poplar_compiler.cc
@@ -1965,7 +1965,7 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
           is_scalar_elementwise_graph,
           /*loaded_from_cache=*/false, std::move(remaped_output),
           std::move(resources.annotations.stream_infos),
-          std::move(resources.annotations.stream_meta_infos),
+          std::move(resources.annotations.host_function_infos),
           PoplarExecutableInfo{
               num_IPUs,
               target_type,
diff --git a/tensorflow/compiler/plugin/poplar/driver/poplar_executable.cc b/tensorflow/compiler/plugin/poplar/driver/poplar_executable.cc
@@ -42,7 +42,7 @@ PoplarExecutableCore::PoplarExecutableCore(
     std::vector<std::vector<Literal>> constant_literal_output,
     bool is_remap_graph, bool is_scalar_elementwise_graph,
     bool loaded_from_cache, std::vector<uint64> remaped_output,
-    StreamInfos&& stream_infos, StreamMetaInfos&& stream_meta_info,
+    StreamInfos&& stream_infos, HostFunctionInfos&& host_function_infos,
     PoplarExecutableInfo&& info)
     : poplar_engine_(std::move(engine)),
       input_output_aliasing_map_(std::move(input_output_aliasing_map)),
@@ -53,7 +53,7 @@ PoplarExecutableCore::PoplarExecutableCore(
       loaded_from_cache_(loaded_from_cache),
       remaped_output_(std::move(remaped_output)),
       stream_infos_(std::move(stream_infos)),
-      stream_meta_infos_(std::move(stream_meta_info)),
+      host_function_infos_(std::move(host_function_infos)),
       info_(std::move(info)) {
   TENSORFLOW_TRACEPOINT();
   PopulateCollectiveBalanceReorderHostRerrangements();
@@ -421,7 +421,7 @@ PoplarExecutableCore::Deserialize(
           /*is_remap_graph=*/false,
           /*is_scalar_elementwise_graph=*/false,
           /*loaded_from_cache=*/true, std::vector<uint64>{}, StreamInfos{},
-          StreamMetaInfos{}, std::move(info));
+          HostFunctionInfos{}, std::move(info));
 
   return executable_core;
 }
diff --git a/tensorflow/compiler/plugin/poplar/driver/poplar_executable.h b/tensorflow/compiler/plugin/poplar/driver/poplar_executable.h
@@ -81,7 +81,7 @@ class PoplarExecutableCore {
       std::vector<std::vector<Literal>> constant_literal_output,
       bool is_remap_graph, bool is_scalar_elementwise_graph,
       bool loaded_from_cache, std::vector<uint64> remaped_output,
-      StreamInfos&& stream_infos, StreamMetaInfos&& stream_meta_info,
+      StreamInfos&& stream_infos, HostFunctionInfos&& host_function_infos,
       PoplarExecutableInfo&& info);
 
   ~PoplarExecutableCore();
@@ -120,6 +120,10 @@ class PoplarExecutableCore {
     return info_.remote_parameter_infos;
   }
 
+  const HostFunctionInfos& GetHostFunctionInfos() const {
+    return host_function_infos_;
+  }
+
   const RemoteParameterHostRearrangements&
   GetRemoteParameterHostRearrangements() const {
     return info_.remote_parameter_host_rearrangements;
@@ -133,10 +137,6 @@ class PoplarExecutableCore {
 
   const StreamInfos& GetStreamInfos() const { return stream_infos_; }
 
-  const StreamMetaInfos& GetStreamMetaInfos() const {
-    return stream_meta_infos_;
-  }
-
   const SendRecvInfos& GetSendInfos() const { return info_.send_infos; }
 
   const SendRecvInfos& GetRecvInfos() const { return info_.recv_infos; }
@@ -202,7 +202,7 @@ class PoplarExecutableCore {
 
   // User op info that is not serialized.
   StreamInfos stream_infos_;
-  StreamMetaInfos stream_meta_infos_;
+  HostFunctionInfos host_function_infos_;
 
   // All the other info that is serialized.
   PoplarExecutableInfo info_;
@@ -273,12 +273,12 @@ class PoplarExecutable : public Executable {
     return executable_core_->GetRemoteParameterInfos();
   }
 
-  const StreamInfos& GetStreamInfos() const {
-    return executable_core_->GetStreamInfos();
+  const HostFunctionInfos& GetHostFunctionInfos() const {
+    return executable_core_->GetHostFunctionInfos();
   }
 
-  const StreamMetaInfos& GetStreamMetaInfos() const {
-    return executable_core_->GetStreamMetaInfos();
+  const StreamInfos& GetStreamInfos() const {
+    return executable_core_->GetStreamInfos();
   }
 
   const SendRecvInfos& GetSendInfos() const {
diff --git a/tensorflow/compiler/plugin/poplar/driver/poplar_executor.cc b/tensorflow/compiler/plugin/poplar/driver/poplar_executor.cc
