Fix weight bitpacking which could lead to non-deterministic behaviour (#377) (#378)

Author: lgeiger

larq_compute_engine/mlir/tests/optimize.mlir

@@ -150,6 +150,7 @@ func @bitpack_bconv2d_filters(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16xf
   %0 = "tf.LceBconv2d"(%arg0, %cst, %arg1, %arg2) {activation = "NONE", channels_in = 3 : i32, filter_format = "OHWI", padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
   return %0 : tensor<256x30x30x16xf32>
 
+  // CHECK: %cst = constant dense<0> : tensor<16x3x3x1xi32>
   // CHECK: %0 = "tf.LceBconv2d"(%arg0, %cst, %arg1, %arg2) {activation = "NONE", channels_in = 3 : i32, data_format = "NHWC", dilations = [1, 1, 1, 1], filter_format = "OHWI_PACKED", pad_values = 0 : i32, padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<256x32x32x3xf32>, tensor<16x3x3x1xi32>, tensor<16xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
   // CHECK-NEXT: return %0
 }

larq_compute_engine/mlir/transforms/optimize.cc

@@ -36,8 +36,9 @@ bool IsConstantValue(Attribute values, float expected_value) {
 
 bool IsConv2DFilter(Attribute filter) {
   if (!filter.isa<DenseElementsAttr>()) return false;
-  if (filter.getType().cast<ShapedType>().getShape().size() != 4) return false;
-  return true;
+  auto filter_type = filter.getType().cast<ShapedType>();
+  return filter_type.getElementType().isF32() &&
+         filter_type.getShape().size() == 4;
 }
 
 DenseElementsAttr Bitpack(PatternRewriter& builder, Attribute x) {
@@ -53,11 +54,17 @@ DenseElementsAttr Bitpack(PatternRewriter& builder, Attribute x) {
   int packed_channels = (unpacked_channels + bitwidth - 1) / bitwidth;
 
   std::vector<PackedType> new_values(num_rows * packed_channels);
+  std::vector<float> old_values(num_rows * unpacked_channels);
+
+  int i = 0;
+  for (float x : dense_elements_iter) {
+    old_values[i++] = x;
+  }
+  assert(i == num_rows * unpacked_channels);
 
-  const float* in_ptr = &(*dense_elements_iter.begin());
   using namespace compute_engine::core;
-  packbits_matrix<BitpackOrder::Canonical>(in_ptr, num_rows, unpacked_channels,
-                                           new_values.data());
+  packbits_matrix<BitpackOrder::Canonical>(
+      old_values.data(), num_rows, unpacked_channels, new_values.data());
 
   RankedTensorType out_tensor_type =
       RankedTensorType::get({shape[0], shape[1], shape[2], packed_channels},