int tensor_input_drain()

in torchaudio/csrc/sox/effects_chain.cpp [32:90]

• 48 lines of code
• 8 McCabe index (conditional complexity)

int tensor_input_drain(sox_effect_t* effp, sox_sample_t* obuf, size_t* osamp) {
  // Retrieve the input Tensor and current index
  auto priv = static_cast<TensorInputPriv*>(effp->priv);
  auto index = priv->index;
  auto tensor = *(priv->waveform);
  auto num_channels = effp->out_signal.channels;

  // Adjust the number of samples to read
  const size_t num_samples = tensor.numel();
  if (index + *osamp > num_samples) {
    *osamp = num_samples - index;
  }
  // Ensure that it's a multiple of the number of channels
  *osamp -= *osamp % num_channels;

  // Slice the input Tensor
  auto chunk = [&]() {
    auto i_frame = index / num_channels;
    auto num_frames = *osamp / num_channels;
    auto t = (priv->channels_first)
        ? tensor.index({Slice(), Slice(i_frame, i_frame + num_frames)}).t()
        : tensor.index({Slice(i_frame, i_frame + num_frames), Slice()});
    return t.reshape({-1});
  }();

  // Convert to sox_sample_t (int32_t)
  switch (chunk.dtype().toScalarType()) {
    case c10::ScalarType::Float: {
      // Need to convert to 64-bit precision so that
      // values around INT32_MIN/MAX are handled correctly.
      chunk = chunk.to(c10::ScalarType::Double);
      chunk *= 2147483648.;
      chunk.clamp_(INT32_MIN, INT32_MAX);
      chunk = chunk.to(c10::ScalarType::Int);
      break;
    }
    case c10::ScalarType::Int: {
      break;
    }
    case c10::ScalarType::Short: {
      chunk = chunk.to(c10::ScalarType::Int);
      chunk *= 65536;
      break;
    }
    case c10::ScalarType::Byte: {
      chunk = chunk.to(c10::ScalarType::Int);
      chunk -= 128;
      chunk *= 16777216;
      break;
    }
    default:
      throw std::runtime_error("Unexpected dtype.");
  }
  // Write to buffer
  chunk = chunk.contiguous();
  memcpy(obuf, chunk.data_ptr<int32_t>(), *osamp * 4);
  priv->index += *osamp;
  return (priv->index == num_samples) ? SOX_EOF : SOX_SUCCESS;
}