[Placeholder] Add support for the runBatch interface.

Add support for the runBatch interface. This commit implements the parts
of runBatch that copy parts of the tensors into the placeholder. The
implementation is zero-copy because we use tensor views.  I hope that
this is a temporary solution and that we'll get rid of runBatch soon
(see #1587).

Author: nadavrot

lib/ExecutionEngine/ExecutionEngine.cpp

@@ -113,8 +113,15 @@ void ExecutionEngine::runBatch(size_t iterations,
 void ExecutionEngine::updateInputsAndRunNetwork(llvm::ArrayRef<Storage *> vars,
                                                 llvm::ArrayRef<Tensor *> inputs,
                                                 size_t sampleIdx) {
-
   llvm::SmallVector<Placeholder *, 8> placeholders;
+
+  // This container saves the tensor slices that were extracted from the inputs.
+  // The tensors are alive during the lifetime of this function. These tensors
+  // must not move, and we must pass them as ArrayRef, so we allocate them here
+  // manually.
+  llvm::SmallVector<Tensor, 8> tempTensorViews;
+  // We need to pass the tensors as array ref, so we reference the static
+  // storage above, that's frozen during the execution of the program.
   llvm::SmallVector<Tensor *, 8> tensors;
 
   // Update the input variables.
@@ -129,7 +136,26 @@ void ExecutionEngine::updateInputsAndRunNetwork(llvm::ArrayRef<Storage *> vars,
     // This is a placeholder that we need to make concrete during the execution
     // of the program.
     placeholders.push_back(cast<Placeholder>(vars[i]));
-    tensors.push_back(inputs[i]);
+
+    // Extract a tensor view from the input values. This has the same
+    // functionality as the logic in loadValueFromTensorSlice that copies the
+    // content of the tensor.
+    auto batchSize = inputs[i]->dims()[0];
+    auto startIdx = sampleIdx % batchSize;
+    auto phDims = placeholders[i]->dims();
+
+    // The start offset is all zeros except for the batch dimension.
+    std::vector<size_t> sliceOffset(inputs[i]->dims().size(), 0);
+    sliceOffset[0] = startIdx;
+
+    // Create the tensor view.
+    tempTensorViews.push_back(inputs[i]->getUnowned(phDims, sliceOffset));
+  }
+
+  // Save the pointer to the allocated tensor view, because the interface
+  // requires ArrayRef of tensors.
+  for (int i = 0, e = tempTensorViews.size(); i < e; i++) {
+    tensors.push_back(&tempTensorViews[i]);
   }
 
   // Run the network.

tests/unittests/BackendTest.cpp

@@ -173,6 +173,8 @@ TEST_P(BackendTest, decoupleCodegenFromGraph) {
   EXPECT_NEAR(HX.at({2}), 9, 1E-5);
 }
 
+/// Check that we can pass information to the execution engine using Placeholder
+/// variables and read it back using Save nodes (in variables).
 TEST(Placeholder, simplePlaceholderValue) {
   Tensor data{99.0, 35.0, 2.0, 3.0};
 
@@ -188,11 +190,42 @@ TEST(Placeholder, simplePlaceholderValue) {
   EE.run({input}, {&data});
 
   auto &res = S->getVariable()->getPayload();
-
-  res.getHandle().dump();
   EXPECT_TRUE(res.isEqual(data));
 }
 
+/// Check that we can pass information to the execution engine using Placeholder
+/// variables and the runBatch API.
+TEST(Placeholder, runBatchTest) {
+  // The input contains two slices of 4 floats each.
+  Tensor data(ElemKind::FloatTy, {4, 4});
+  // Fill the array with the pattern: [0 1 2 3; 10, 11, 12, 13; 20 21 22 23 ...]
+  for (size_t i = 0; i < 4; i++) {
+    for (size_t j = 0; j < 4; j++) {
+      data.getHandle().at({i, j}) = i * 10 + j;
+    }
+  }
+
+  ExecutionEngine EE{BackendKind::Interpreter};
+  auto &mod = EE.getModule();
+
+  Function *F = mod.createFunction("main");
+  auto *input = mod.createPlaceholder(ElemKind::FloatTy, {1, 4}, "input");
+  SaveNode *S = F->createSave("ret", input);
+
+  EE.compile(CompilationMode::Infer, F);
+
+  // Run the batch for 2 iterations:
+  EE.runBatch(2, {input}, {&data});
+
+  Tensor expected{10, 11, 12, 13};
+  auto EH = expected.getHandle();
+  auto RH = S->getVariable()->getPayload().getHandle();
+
+  for (size_t i = 0; i < 4; i++) {
+    EXPECT_EQ(RH.raw(i), EH.raw(i));
+  }
+}
+
 INSTANTIATE_TEST_CASE_P(Interpreter, BackendTest,
                         ::testing::Values(BackendKind::Interpreter));