Simplified error handling for bit_coder and run_length_gamma_kernels.

PiperOrigin-RevId: 470074035
Change-Id: I9452966aef1f56a5b0c7570b24516ee60bd2b64c

Author: Johannes Ballé

tensorflow_compression/cc/kernels/run_length_gamma_kernels.cc

@@ -16,6 +16,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <array>
+#include <cassert>
 #include <cmath>
 #include <cstdint>
 #include <cstring>
@@ -28,17 +29,16 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_types.h"
-#include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow_compression/cc/lib/bit_coder.h"
 
 namespace tensorflow_compression {
 namespace {
 namespace errors = tensorflow::errors;
 using tensorflow::DEVICE_CPU;
+using tensorflow::FromAbslStatus;
 using tensorflow::OpKernel;
 using tensorflow::OpKernelConstruction;
 using tensorflow::OpKernelContext;
@@ -69,36 +69,38 @@ class RunLengthGammaEncodeOp : public OpKernel {
     // any zeros were present in the input array, then the encoded size would be
     // strictly smaller by kMaxGammaBits and bigger by the difference in
     // encoding (the existing zero run length + 1).
-    BitWriter enc;
-    enc.Allocate(data.size() * (2 + enc.kMaxGammaBits));
+    BitWriter enc(data.size() * (2 + enc.kMaxGammaBits));
     // Save number of zeros + 1 preceding next non-zero element.
     uint32_t zero_ct = 1;
 
     // Iterate through data tensor.
-    for (size_t i = 0; i < data.size(); i++) {
+    for (int64_t i = 0; i < data.size(); i++) {
+      int32_t sample = data(i);
       // Increment zero count.
-      if (data(i) == 0) {
+      if (sample == 0) {
         zero_ct += 1;
       } else {
         // Encode run length of zeros.
         enc.WriteGamma(zero_ct);
         // Encode sign of value.
-        enc.WriteOneBit(data(i) > 0);
+        enc.WriteOneBit(sample > 0);
         // Encode magnitude of value.
-        DCHECK_NE(data(i), std::numeric_limits<int32_t>::min());
-        enc.WriteGamma(std::abs(data(i)));
-        // Reset zero count (1 because Gamma cannot encode 0).
+        if (sample == std::numeric_limits<int32_t>::min()) {
+          // We can't encode int32 minimum. Encode closest value instead.
+          sample += 1;
+        }
+        enc.WriteGamma(std::abs(sample));
+        // Reset zero count (1 because gamma cannot encode 0).
         zero_ct = 1;
       }
     }
     if (zero_ct > 1) {
       enc.WriteGamma(zero_ct);
     }
 
-    // Pad any remaining bits in last byte with 0.
-    enc.ZeroPadToByte();
     // Write encoded bitstring to code.
-    code->assign(enc.GetData(), enc.GetBytesWritten());
+    auto encoded = enc.GetData();
+    code->assign(encoded.data(), encoded.size());
   }
 };
 
@@ -137,32 +139,33 @@ class RunLengthGammaDecodeOp : public OpKernel {
     // Fill data tensor with zeros.
     std::memset(data.data(), 0, data.size() * sizeof(data(0)));
 
-    for (size_t i = 0; i < data.size(); i++) {
+    for (int64_t i = 0; i < data.size(); i++) {
       // Get number of zeros.
-      uint32_t num_zeros = dec.ReadGamma();
+      auto num_zeros = dec.ReadGamma();
+      OP_REQUIRES(context, num_zeros.ok(), FromAbslStatus(num_zeros.status()));
+
       // Advance the index to the next non-zero element.
-      i += num_zeros - 1;
+      i += *num_zeros - 1;
 
       // Account for case where the last element is zero.
-      if (i == data.size()) {
+      // Check if past the last element.
+      if (i >= data.size()) {
+        OP_REQUIRES(context, i == data.size(),
+                    errors::DataLoss("Decoded past end of tensor."));
         break;
       }
-      // TODO(nicolemitchell): return error status instead of crashing
-      DCHECK_LT(i, data.size());
 
       // Get sign of value.
-      uint32_t positive = dec.ReadOneBit();
+      auto positive = dec.ReadOneBit();
+      OP_REQUIRES(context, positive.ok(), FromAbslStatus(positive.status()));
 
-      // Get value.
-      uint32_t value = dec.ReadGamma();
+      // Get magnitude.
+      auto magnitude = dec.ReadGamma();
+      OP_REQUIRES(context, magnitude.ok(), FromAbslStatus(magnitude.status()));
 
       // Write value to data tensor element at index.
-      DCHECK_LE(value, std::numeric_limits<int32_t>::max());
-      data(i) = positive ? value : -static_cast<int32_t>(value);
+      data(i) = *positive ? *magnitude : -*magnitude;
     }
-
-    OP_REQUIRES(context, dec.Close().ok(),
-                tensorflow::errors::DataLoss("Decoding error."));
   }
 };
 

tensorflow_compression/cc/kernels/run_length_gamma_kernels_test.cc

@@ -35,13 +35,9 @@ limitations under the License.
 #include "tensorflow/core/graph/node_builder.h"
 #include "tensorflow/core/graph/testlib.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
-#include "tensorflow/core/lib/core/bits.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/lib/random/simple_philox.h"
 #include "tensorflow/core/platform/stacktrace_handler.h"
 #include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/public/session.h"
-#include "tensorflow/core/public/session_options.h"
 #include "tensorflow_compression/cc/lib/bit_coder.h"
 
 namespace tensorflow_compression {
@@ -198,36 +194,35 @@ TEST_F(BitCodingOpsTest, ManualEncodeWithBitcodingLibrary) {
   TF_ASSERT_OK(RunEncodeOp({data_tensor}, &code_tensor));
 
   // Use bitcoding library to encode data.
-  BitWriter enc_ = BitWriter();
-  enc_.Allocate(16);
+  BitWriter enc_ = BitWriter(16);
   enc_.WriteGamma(2);   // one zero
   enc_.WriteOneBit(0);  // negative
   enc_.WriteGamma(3);   // 3
   enc_.WriteGamma(1);   // no zeros
   enc_.WriteOneBit(1);  // positive
   enc_.WriteGamma(1);   // 1
-  enc_.ZeroPadToByte();
   Tensor expected_code_tensor(DT_STRING, {});
-  expected_code_tensor.scalar<tstring>()().assign(enc_.GetData(), 2);
+  auto encoded = enc_.GetData();
+  expected_code_tensor.scalar<tstring>()().assign(encoded.data(),
+                                                  encoded.size());
 
   // Check that code_tensor has expected value.
   test::ExpectTensorEqual<tstring>(code_tensor, expected_code_tensor);
 }
 
 TEST_F(BitCodingOpsTest, ManualDecodeWithBitcodingLibrary) {
   // Use bitcoding library to manually encode [-3, 1, 0, 0] into code.
-  BitWriter enc_ = BitWriter();
-  enc_.Allocate(16);
+  BitWriter enc_ = BitWriter(16);
   enc_.WriteGamma(1);   // no zeros
   enc_.WriteOneBit(0);  // negative
   enc_.WriteGamma(3);   // 3
   enc_.WriteGamma(1);   // no zeros
   enc_.WriteOneBit(1);  // positive
   enc_.WriteGamma(1);   // 1
   enc_.WriteGamma(3);   // two zeros
-  enc_.ZeroPadToByte();
   Tensor code_tensor(DT_STRING, {});
-  code_tensor.scalar<tstring>()().assign(enc_.GetData(), 2);
+  auto encoded = enc_.GetData();
+  code_tensor.scalar<tstring>()().assign(encoded.data(), encoded.size());
 
   Tensor shape_tensor(DT_INT32, {1});
   shape_tensor.flat<int32_t>().setValues({4});
@@ -242,7 +237,6 @@ TEST_F(BitCodingOpsTest, ManualDecodeWithBitcodingLibrary) {
   test::ExpectTensorEqual<int32_t>(data_tensor, expected_data_tensor);
 }
 
-// TODO(nicolemitchell) Strengthen these consistency checks.
 TEST_F(BitCodingOpsTest, EncodeConsistent) {
   Tensor data_tensor(DT_INT32, {4});
   data_tensor.flat<int32_t>().setValues({-6, 3, 0, 0});
@@ -259,7 +253,6 @@ TEST_F(BitCodingOpsTest, EncodeConsistent) {
 }
 
 TEST_F(BitCodingOpsTest, DecodeConsistent) {
-  // Manually encode some data into code.
   char code[] = {0b11010001, 0b01101101};  // [-6, 3, 0, 0]
 
   Tensor code_tensor(DT_STRING, {});
@@ -277,6 +270,14 @@ TEST_F(BitCodingOpsTest, DecodeConsistent) {
   // Check that decoded data has expected values.
   test::ExpectTensorEqual<int32_t>(data_tensor, expected_data_tensor);
 }
+
+// TODO(nicolemitchell,jonycgn) Add more corner cases to unit tests.
+// Examples: decode empty string (null pointer), decode strings that end
+// prematurely, decode long string of zeros that causes overflow in ReadGamma,
+// decode incorrect run length that exceeds tensor size, encode int32::min
+// tensor, encode tensor with very large values to ensure it doesn't exceed
+// allocated buffer, encode gamma values <= 0, ...
+
 }  // namespace
 }  // namespace tensorflow_compression