Function SetupAndRun

dali/python/backend_impl.cc:2932–3029

}

void SetupAndRun(OperatorBase &self, Workspace &ws, std::optional<int> batch_size) {
  DomainTimeRange setup_and_run_tr("SetupAndRun " + GetOpDisplayName(self.GetSpec(), true),
                                   kDynamicDefaultColor);
  std::vector<dali::OutputDesc> out_descs;
  const auto &spec = self.GetSpec();
  if (ws.NumOutput() != 0)
    throw std::runtime_error("Workspace already has outputs defined");

  auto AdjustLayout = [&](int idx, const TensorLayout &layout, auto backend) {
    using Backend = decltype(backend);
    const auto &inp = ws.Input<Backend>(idx);
    auto tl = std::make_shared<TensorList<Backend>>();
    tl->ShareData(inp);
    tl->SetLayout(layout);
    ws.SetInput(idx, std::move(tl));
  };

  auto &schema = spec.GetSchemaOrDefault();
  for (int i = 0; i < spec.NumRegularInput(); i++) {
    auto layout = ws.GetInputLayout(i);
    auto ndim = ws.GetInputDim(i);
    auto adjusted_layout = schema.GetInputLayout(i, ndim, layout);
    if (layout != adjusted_layout) {
      if (ws.InputIsType<GPUBackend>(i)) {
        AdjustLayout(i, adjusted_layout, GPUBackend());
      } else {
        AdjustLayout(i, adjusted_layout, CPUBackend());
      }
    }
  }

  bool pinned = ws.has_stream();
  std::optional<int> dev_id;
  if (ws.output_order().is_device())
    dev_id = ws.output_order().device_id();

  for (int i = 0; i < spec.NumOutput(); i++) {
    if (spec.OutputDevice(i) == StorageDevice::CPU) {
      auto out = std::make_shared<TensorList<CPUBackend>>();
      if (dev_id.has_value())
        out->set_device_id(*dev_id);
      out->set_order(AccessOrder::host(), false);  // make sure outputs are allocated in host order
      out->set_pinned(pinned);
      ws.AddOutput(std::move(out));
    } else {
      auto out = std::make_shared<TensorList<GPUBackend>>();
      if (!dev_id.has_value()) {
        throw std::logic_error(
          "Internal error: Workspace for a GPU operator doesn't have a valid device_id.");
      }
      out->set_device_id(*dev_id);
      out->set_order(ws.output_order(), false);
      ws.AddOutput(std::move(out));
    }
  }

  if (batch_size.has_value()) {
    ws.SetBatchSizes(batch_size.value());