Merge master again (#2097, #2117)

include/ctranslate2/storage_view.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <limits>
 #include <ostream>
 #include <vector>
 
@@ -29,9 +30,25 @@ namespace ctranslate2 {
   using Shape = std::vector<dim_t>;
 
   inline dim_t compute_size(const Shape& shape) {
+    bool has_zero_dim = false;
+    for (const dim_t dim : shape) {
+      if (dim < 0)
+        THROW_INVALID_ARGUMENT("shape contains a negative dimension");
+      if (dim == 0)
+        has_zero_dim = true;
+    }
+
+    // A zero dimension makes the size 0 whatever its position and whatever the
+    // other dimensions, so nothing is multiplied and no overflow is possible.
+    if (has_zero_dim)
+      return 0;
+
     dim_t size = 1;
-    for (const dim_t dim : shape)
+    for (const dim_t dim : shape) {
+      if (size > std::numeric_limits<dim_t>::max() / dim)
+        THROW_INVALID_ARGUMENT("shape leads to an overflowing tensor size");
       size *= dim;
+    }
     return size;
   }
 

python/ctranslate2/specs/model_spec.py

@@ -233,6 +233,8 @@ def _quantize(spec, name, value):
                         value = value.reshape(value.shape[0], -1)
                     amax = np.amax(np.absolute(value), axis=1)
                     amax[amax == 0] = 127.0
+                    # Keep the scale representable in float32 for tiny nonzero rows.
+                    amax = np.maximum(amax, np.float32(127) / np.finfo(np.float32).max)
                     scale = 127.0 / amax
                     value *= np.expand_dims(scale, 1)
                     value = np.rint(value)

python/tests/test_spec.py

@@ -116,6 +116,21 @@ def __init__(self):
     )
 
 
+def test_int8_quantization_tiny_weights():
+    class Spec(ctranslate2.specs.LayerSpec):
+        def __init__(self):
+            self.weight = np.array([[1e-38, 0, 0, 0]], dtype=np.float32)
+            self.weight_scale = OPTIONAL
+
+    spec = Spec()
+    spec.validate()
+    with np.errstate(over="raise", invalid="raise"):
+        spec.optimize(quantization="int8")
+    assert np.isfinite(spec.weight.numpy()).all()
+    assert np.isfinite(spec.weight_scale.numpy()).all()
+    assert (spec.weight_scale.numpy() > 0).all()
+
+
 @pytest.mark.parametrize(
     "quantization,expected_weight,expected_weight_scale,expected_bias",
     [

src/ops/conv1d_cpu.cc

@@ -36,7 +36,7 @@ namespace ctranslate2 {
                               const_cast<void*>(weight.buffer()));
 
       dnnl::memory::dims stride{_stride};
-      dnnl::memory::dims dilation{_dilation > 1 ? _dilation : 0};
+      dnnl::memory::dims dilation{_dilation - 1};
       dnnl::memory::dims padding{_padding};
 
       std::unique_ptr<dnnl::convolution_forward::primitive_desc> conv_pd;
@@ -130,9 +130,6 @@ namespace ctranslate2 {
     void
     Conv1D::compute<Device::CPU, float>(const StorageView &input, const StorageView &weight, const StorageView *bias,
                                         StorageView &output, const StorageView *qscale) const {
-      if (_dilation != 1)
-        throw std::runtime_error("Dilation is not supported in this Conv1D implementation");
-
       compute_with_gemm(input, weight, output, qscale);
       // Add bias
       apply_bias_and_activation(output, bias, _activation_type, /*residual=*/nullptr, /*axis=*/-2);
@@ -218,6 +215,7 @@ namespace ctranslate2 {
       const dim_t batch_size = input.dim(0);
       const dim_t in_channels = input.dim(1);
       const dim_t input_length = input.dim(2);
+      const dim_t effective_kernel_size = _dilation * (kernel_size - 1) + 1;
       auto* out = output.data <float>();
       const auto* in = input.data <float>();
       dim_t out_offset = 0;
@@ -226,11 +224,11 @@ namespace ctranslate2 {
       const dim_t in_group_stride = in_channels_per_group * input_length;
       for (dim_t batch_offset = 0; batch_offset < batch_size * in_batch_stride; batch_offset += in_batch_stride) {
         for (dim_t group_offset = batch_offset; group_offset < (batch_offset + _groups * in_group_stride); group_offset += in_group_stride) {
-          for (dim_t ti = -_padding; ti <= (input_length - kernel_size + _padding); ti += _stride) {
+          for (dim_t ti = -_padding; ti <= (input_length - effective_kernel_size + _padding); ti += _stride) {
             for (dim_t c = group_offset; c < (group_offset + in_channels_per_group * input_length); c += input_length) {
               for (int k = 0; k < kernel_size; k++) {
                 // Fill items in [0, input_length) range
-                auto window_i = k + ti;
+                auto window_i = k * _dilation + ti;
                 if (0 <= window_i && window_i < input_length) {
                   out[out_offset] = in[window_i + c];
                 }

src/storage_view.cc

@@ -190,12 +190,15 @@ namespace ctranslate2 {
   StorageView& StorageView::reshape(Shape new_shape) {
     dim_t unknown_dim = -1;
     dim_t known_size = 1;
+    bool has_zero_dim = false;
 
     for (size_t i = 0; i < new_shape.size(); ++i) {
       const dim_t dim = new_shape[i];
 
       if (dim >= 0) {
-        known_size *= dim;
+        // the empty dimensions are only recorded here: the size is computed once
+        // the whole shape is known to contain no negative dimension
+        has_zero_dim = has_zero_dim || dim == 0;
       } else if (dim == -1) {
         if (unknown_dim >= 0)
           THROW_INVALID_ARGUMENT("only one dimension can be set to -1, got -1 for dimensions "
@@ -207,6 +210,25 @@ namespace ctranslate2 {
       }
     }
 
+    // A zero dimension makes the new size 0 whatever its position and whatever the
+    // other dimensions, so nothing is multiplied and no overflow is possible. The
+    // unknown dimension can't be recovered from an empty size, and the compatibility
+    // check below would take the modulo by 0.
+    if (has_zero_dim) {
+      if (unknown_dim >= 0)
+        THROW_INVALID_ARGUMENT("can't infer dimension " + std::to_string(unknown_dim)
+                               + " of a new shape with a zero dimension");
+      known_size = 0;
+    } else {
+      for (const dim_t dim : new_shape) {
+        if (dim <= 0)
+          continue;
+        if (known_size > std::numeric_limits<dim_t>::max() / dim)
+          THROW_INVALID_ARGUMENT("new shape leads to an overflowing tensor size");
+        known_size *= dim;
+      }
+    }
+
     if (unknown_dim >= 0) {
       if (_size % known_size != 0)
         THROW_INVALID_ARGUMENT("current size (" + std::to_string(_size)
@@ -261,6 +283,10 @@ namespace ctranslate2 {
 
   StorageView& StorageView::grow(dim_t dim, dim_t size) {
     GUARD_DIM(dim, rank());
+    if (size < 0)
+      THROW_INVALID_ARGUMENT("can't grow a dimension by a negative size");
+    if (_shape[dim] > std::numeric_limits<dim_t>::max() - size)
+      THROW_INVALID_ARGUMENT("new dimension size leads to an overflowing tensor size");
     return resize(dim, _shape[dim] + size);
   }
 

tests/ops_test.cc

@@ -1452,10 +1452,24 @@ TEST_P(OpDeviceFPTest, Conv1DPaddingAndStride) {
   expect_storage_eq(output.to_float32(), expected, error);
 }
 
+TEST(OpTest, Conv1DDilationCPU) {
+  const StorageView input({1, 1, 7}, std::vector<float>{1, 2, 3, 4, 5, 6, 7});
+  const StorageView weight({1, 1, 3}, std::vector<float>{1, 2, 3});
+  const StorageView bias({1}, std::vector<float>{1});
+  StorageView output;
+
+  ops::Conv1D(1, 0, 2)(input, weight, output);
+  expect_storage_eq(output, StorageView({1, 1, 3}, std::vector<float>{22, 28, 34}));
+
+  ops::Conv1D(1, 0, 3)(input, weight, output);
+  expect_storage_eq(output, StorageView({1, 1, 1}, std::vector<float>{30}));
+
+  ops::Conv1D(2, 2, 2)(input, weight, bias, output);
+  expect_storage_eq(output, StorageView({1, 1, 4}, std::vector<float>{12, 23, 35, 20}));
+}
+
 TEST_P(OpDeviceFPTest, Conv1DDilation) {
   const Device device = GetParam().device;
-  if (device == Device::CPU)
-      GTEST_SKIP() << "Dilated convolution is not implemented for CPU.";
   const DataType dtype = GetParam().dtype;
   const float error = GetParam().error;
   const StorageView expected({2, 4, 1}, std::vector<float>{

tests/storage_view_test.cc

@@ -148,3 +148,114 @@ INSTANTIATE_TEST_SUITE_P(CUDA, StorageViewDeviceTest, ::testing::Values(Device::
 #ifdef CT2_WITH_MPS
 INSTANTIATE_TEST_SUITE_P(MPS, StorageViewDeviceTest, ::testing::Values(Device::MPS));
 #endif
+
+TEST(StorageViewTest, OverflowingShapeThrows) {
+  // (2^62 + 1) * 8 overflows dim_t: the shape must be rejected, not wrapped
+  // to a small size that passes allocation.
+  EXPECT_THROW(StorageView(Shape{(int64_t{1} << 62) + 1, 8}, DataType::INT32),
+               std::invalid_argument);
+}
+
+TEST(StorageViewTest, OverflowingReshapeThrows) {
+  StorageView view(Shape{8}, DataType::INT32);
+  EXPECT_THROW(view.reshape(Shape{(int64_t{1} << 62) + 1, 8}), std::invalid_argument);
+}
+
+TEST(StorageViewTest, ZeroDimIsValidWhateverItsPosition) {
+  // The size of a shape with a zero dimension is 0, so the other dimensions are
+  // never multiplied and the overflowing prefix below is not an error. Multiplying
+  // before looking for the zero dimension made the result depend on its position.
+  EXPECT_THROW(compute_size(Shape{(int64_t{1} << 62) + 1, 8}), std::invalid_argument);
+  EXPECT_EQ(compute_size(Shape{(int64_t{1} << 62) + 1, 8, 0}), 0);
+  EXPECT_EQ(compute_size(Shape{0, (int64_t{1} << 62) + 1, 8}), 0);
+
+  StorageView view(Shape{(int64_t{1} << 62) + 1, 8, 0}, DataType::INT32);
+  EXPECT_EQ(view.size(), 0);
+  EXPECT_EQ(view.dim(2), 0);
+  view.reshape(Shape{0, (int64_t{1} << 62) + 1, 8});
+  EXPECT_EQ(view.size(), 0);
+  EXPECT_EQ(view.dim(0), 0);
+}
+
+TEST(StorageViewTest, ZeroDimDoesNotSkipNegativeCheck) {
+  EXPECT_THROW(compute_size(Shape{0, -2}), std::invalid_argument);
+  EXPECT_THROW(compute_size(Shape{(int64_t{1} << 62) + 1, 0, -1}), std::invalid_argument);
+
+  StorageView view(Shape{0}, DataType::INT32);
+  EXPECT_THROW(view.reshape(Shape{0, 8, -3}), std::invalid_argument);
+}
+
+TEST(StorageViewTest, UnknownDimWithZeroDimThrows) {
+  // An unknown dimension can't be inferred from an empty shape, and the known size
+  // of 0 used to be the divisor of the compatibility check below.
+  StorageView view(Shape{8}, DataType::INT32);
+  EXPECT_THROW(view.reshape(Shape{0, -1}), std::invalid_argument);
+  EXPECT_THROW(view.reshape(Shape{-1, 0}), std::invalid_argument);
+}
+
+TEST(StorageViewTest, GrowRejectsNegativeIncrement) {
+  // grow() with a negative size used to reach resize() and silently shrink the
+  // dimension, and INT64_MIN made the addition wrap.
+  StorageView view(Shape{4}, DataType::INT32);
+  for (const dim_t size : {dim_t(-1), std::numeric_limits<dim_t>::min()}) {
+    try {
+      view.grow(0, size);
+      FAIL() << "Expected a negative growth to be rejected";
+    } catch (const std::invalid_argument& e) {
+      EXPECT_NE(std::string(e.what()).find("can't grow a dimension by a negative size"),
+                std::string::npos);
+    }
+  }
+
+  // A rejected grow leaves the storage untouched.
+  EXPECT_EQ(view.dim(0), 4);
+  EXPECT_EQ(view.size(), 4);
+}
+
+TEST(StorageViewTest, GrowRejectsAdditionOverflow) {
+  // The zero dimension makes the size 0, so this overflowing first dimension costs
+  // no memory and still exercises the guard on the addition itself.
+  StorageView view(Shape{std::numeric_limits<dim_t>::max(), 0}, DataType::INT32);
+  EXPECT_EQ(view.size(), 0);
+
+  EXPECT_THROW(view.grow(0, 1), std::invalid_argument);
+  EXPECT_THROW(view.grow(0, std::numeric_limits<dim_t>::max()), std::invalid_argument);
+
+  // Growing the zero dimension by INT64_MAX doesn't overflow the addition (0 +
+  // INT64_MAX), so grow() lets it through and compute_size() rejects the resulting
+  // shape instead of wrapping it.
+  EXPECT_THROW(view.grow(1, std::numeric_limits<dim_t>::max()), std::invalid_argument);
+
+  EXPECT_EQ(view.dim(0), std::numeric_limits<dim_t>::max());
+  EXPECT_EQ(view.dim(1), 0);
+  EXPECT_EQ(view.size(), 0);
+}
+
+TEST(StorageViewTest, GrowingToAnUnallocatableSizeIsRejected) {
+  // INT64_MAX is a valid dimension and INT64_MAX elements still fit in a dim_t, so
+  // neither grow() nor compute_size() rejects it: reserve() catches the byte size
+  // before allocating. The three guards leave no path through to the allocator.
+  StorageView view(Shape{0}, DataType::INT32);
+  EXPECT_THROW(view.grow(0, std::numeric_limits<dim_t>::max()), std::runtime_error);
+  EXPECT_EQ(view.size(), 0);
+}
+
+TEST(StorageViewTest, GrowIncreasesShape) {
+  StorageView view(Shape{2, 3}, DataType::INT32);
+  EXPECT_EQ(view.size(), 6);
+
+  view.grow(0, 1);
+  EXPECT_EQ(view.dim(0), 3);
+  EXPECT_EQ(view.dim(1), 3);
+  EXPECT_EQ(view.size(), 9);
+
+  view.grow(1, 2);
+  EXPECT_EQ(view.dim(1), 5);
+  EXPECT_EQ(view.size(), 15);
+
+  // Growing an empty dimension back to a positive size stays valid.
+  StorageView empty_view(Shape{0}, DataType::INT32);
+  empty_view.grow(0, 4);
+  EXPECT_EQ(empty_view.dim(0), 4);
+  EXPECT_EQ(empty_view.size(), 4);
+}