Method RunTest

tensorflow/core/kernels/collective_nccl_test.cc:181–222

  }

  void RunTest(int num_ranks, int input_length, int instance_key) {
    Init(num_ranks, instance_key);
    std::vector<float> expected;
    InitExpected(&expected, input_length, num_ranks);
    if (VLOG_IS_ON(3)) {
      string str_buf;
      for (const auto& x : expected) {
        strings::StrAppend(&str_buf, " ", x);
      }
      VLOG(3) << "Expected output " << str_buf;
    }
    for (int rank = 0; rank < num_ranks; ++rank) {
      DeviceInstance* instance = instances_[rank].get();
      instance->InitTensor(DT_FLOAT, TensorShape({input_length}),
                           [this, rank](Tensor* t) { InitInput(t, rank); });
    }
    RunCollective();
    // Confirm that every rank computed the same correct value.
    for (int rank = 0; rank < instances_.size(); ++rank) {
      TF_ASSERT_OK(instances_[rank]->status_);
      Tensor* output = &instances_[rank]->output_;
      const int output_length = output->NumElements();
      VLOG(2) << "rank " << rank << " output " << output << " buf "
              << DMAHelper::base(output);
      Tensor actual(DT_FLOAT, TensorShape({output_length}));
      Notification note;
      Device* dev = instances_[rank]->device_;
      auto* dev_info = dev->tensorflow_gpu_device_info();
      dev_info->default_context->CopyDeviceTensorToCPU(
          output, /*tensor_name=*/"", dev, &actual, [&note](const Status& s) {
            TF_CHECK_OK(s);
            note.Notify();
          });
      note.WaitForNotification();
      VLOG(3) << "rank " << rank << " got output tensor "
              << actual.DebugString(output_length);
      for (int i = 0; i < output_length; ++i) {
        EXPECT_FLOAT_EQ(expected[i], actual.template flat<float>()(i))
            << "Mismatch at rank " << rank << " index " << i;
      }
    }
  }

  std::unique_ptr<OpKernel> GetCollectiveReduceOpKernel(
      const CollectiveParams& params, Tensor* input, DeviceBase* device) {