Replace sox conversion macros with Tensor op

Conversion from signed integer with left shift causes UB (negative left shift)
pytorch · Mar 7, 2021 · fd2e5c4 · fd2e5c4
1 parent 7b85f1c
commit fd2e5c4
Showing 1 changed file with 23 additions and 17 deletions.
diff --git a/torchaudio/csrc/sox/effects_chain.cpp b/torchaudio/csrc/sox/effects_chain.cpp
@@ -69,44 +69,50 @@ int tensor_input_drain(sox_effect_t* effp, sox_sample_t* obuf, size_t* osamp) {
   *osamp -= *osamp % num_channels;
 
   // Slice the input Tensor
-  const auto 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}).contiguous();
+    return t.reshape({-1});
   }();
 
   // Convert to sox_sample_t (int32_t) and write to buffer
-  SOX_SAMPLE_LOCALS;
-  switch (tensor_.dtype().toScalarType()) {
+  switch (chunk.dtype().toScalarType()) {
     case c10::ScalarType::Float: {
-      auto ptr = tensor_.data_ptr<float_t>();
+      // The following code snipet should work, but, there is a bug in
+      // FloatTensor with values near INT32_MAX, so until that bug is resolved,
+      // we use the sox's macro to convert the values.
+      // https://github.com/pytorch/audio/issues/1337
+      //
+      // chunk *= 2147483648.;
+      // chunk.clamp_(INT32_MIN, INT32_MAX);
+      // chunk = chunk.to(c10::ScalarType::Int);
+      // memcpy(obuf, chunk.data_ptr<int32_t>(), *osamp*4);
+      chunk = chunk.contiguous();
+      SOX_SAMPLE_LOCALS;
       for (size_t i = 0; i < *osamp; ++i) {
         obuf[i] = SOX_FLOAT_32BIT_TO_SAMPLE(ptr[i], effp->clips);
       }
       break;
     }
     case c10::ScalarType::Int: {
-      auto ptr = tensor_.data_ptr<int32_t>();
-      for (size_t i = 0; i < *osamp; ++i) {
-        obuf[i] = SOX_SIGNED_32BIT_TO_SAMPLE(ptr[i], effp->clips);
-      }
+      chunk = chunk.contiguous();
+      memcpy(obuf, chunk.data_ptr<int32_t>(), *osamp*4);
       break;
     }
     case c10::ScalarType::Short: {
-      auto ptr = tensor_.data_ptr<int16_t>();
-      for (size_t i = 0; i < *osamp; ++i) {
-        obuf[i] = SOX_SIGNED_16BIT_TO_SAMPLE(ptr[i], effp->clips);
-      }
+      chunk = chunk.to(c10::ScalarType::Int);
+      chunk *= 65536;
+      memcpy(obuf, chunk.data_ptr<int32_t>(), *osamp*4);
       break;
     }
     case c10::ScalarType::Byte: {
-      auto ptr = tensor_.data_ptr<uint8_t>();
-      for (size_t i = 0; i < *osamp; ++i) {
-        obuf[i] = SOX_UNSIGNED_8BIT_TO_SAMPLE(ptr[i], effp->clips);
-      }
+      chunk = chunk.to(c10::ScalarType::Int);
+      chunk -= 128;
+      chunk *= 16777216;
+      memcpy(obuf, chunk.data_ptr<int32_t>(), *osamp*4);
       break;
     }
     default: