[luci] Introduce Compress weights pass

This commit introduces CopressWeightsPass for Conv2D

ONE-DCO-1.0-Signed-off-by: Vyacheslav Bazhenov <[email protected]>

compiler/circle2circle/src/Circle2Circle.cpp

@@ -184,6 +184,8 @@ int entry(int argc, char **argv)
              "This will convert single input Transpose to Reshape");
   add_switch(arser, "--expand_broadcast_const", "This will expand broadcastable constant inputs");
   add_switch(arser, "--unroll_unidirseqlstm", "Unroll UnidirectionalSequenceLSTM operator.");
+  add_switch(arser, "--compress_weights_huffman",
+             "Loseless weights compression with Huffman encoding.");
   add_switch(arser, "--convert_nchw_to_nhwc",
              "Experimental: This will convert NCHW operators to NHWC under the assumption that "
              "input model is NCHW.");
@@ -405,6 +407,8 @@ int entry(int argc, char **argv)
     options->enable(Algorithms::ExpandBroadcastConst);
   if (arser.get<bool>("--unroll_unidirseqlstm"))
     options->enable(Algorithms::UnrollUnidirSeqLSTM);
+  if (arser.get<bool>("--compress_weights_huffman"))
+    options->enable(Algorithms::CompressWeights);
 
   // NOTE Experimental options; these will be removed someday
   //      Add experimental options here

compiler/luci/export/src/CircleBuiltinTypesExtractor.h

@@ -108,10 +108,10 @@ class BuiltinOptionsExtractor final
   // flatbuffers::Offset<void> visit(luci::CircleConst *)
   flatbuffers::Offset<void> visit(luci::CircleConv2D *node)
   {
-    return circle::CreateConv2DOptions(_builder, getOpPadding(node->padding()), node->stride()->w(),
-                                       node->stride()->h(),
-                                       to_circle_actfunc(node->fusedActivationFunction()),
-                                       node->dilation()->w(), node->dilation()->h())
+    return circle::CreateConv2DOptions(
+             _builder, getOpPadding(node->padding()), node->stride()->w(), node->stride()->h(),
+             to_circle_actfunc(node->fusedActivationFunction()), node->dilation()->w(),
+             node->dilation()->h(), to_circle_weightcompressiontype(node->weightCompression()))
       .Union();
   }
   flatbuffers::Offset<void> visit(luci::CircleCos *)

compiler/luci/export/src/CircleExporterUtils.cpp

@@ -24,7 +24,19 @@
 
 namespace luci
 {
-
+circle::WeightCompressionType to_circle_weightcompressiontype(luci::WeightCompression type)
+{
+  switch (type)
+  {
+    case luci::WeightCompression::NONE:
+      return circle::WeightCompressionType_NONE;
+    case luci::WeightCompression::HUFFMAN:
+      return circle::WeightCompressionType_Huffman;
+    default:
+      INTERNAL_EXN_V("trying to convert unsupported luci::WeightCompression",
+                     oops::to_uint32(type));
+  }
+}
 circle::ActivationFunctionType to_circle_actfunc(luci::FusedActFunc func)
 {
   switch (func)

compiler/luci/export/src/CircleExporterUtils.h

@@ -28,7 +28,7 @@
 
 namespace luci
 {
-
+circle::WeightCompressionType to_circle_weightcompressiontype(luci::WeightCompression type);
 circle::ActivationFunctionType to_circle_actfunc(luci::FusedActFunc func);
 circle::TensorType to_circle_tensortype(loco::DataType type);
 circle::MirrorPadMode to_circle_mirrorpadmode(luci::MirrorPadMode mode);

compiler/luci/export/src/CircleTensorExporter.cpp

@@ -559,6 +559,8 @@ bool has_same_values(luci::CircleConst *lhs, luci::CircleConst *rhs)
       return has_same_elements<loco::DataType::S4>(lhs, rhs);
 
     case loco::DataType::S8:
+      if (lhs->size<loco::DataType::S8>() != rhs->size<loco::DataType::S8>())
+        return false;
       return has_same_elements<loco::DataType::S8>(lhs, rhs);
 
     case loco::DataType::S16:
@@ -574,6 +576,8 @@ bool has_same_values(luci::CircleConst *lhs, luci::CircleConst *rhs)
       return has_same_elements<loco::DataType::U4>(lhs, rhs);
 
     case loco::DataType::U8:
+      if (lhs->size<loco::DataType::U8>() != rhs->size<loco::DataType::U8>())
+        return false;
       return has_same_elements<loco::DataType::U8>(lhs, rhs);
 
     case loco::DataType::BOOL:

compiler/luci/lang/include/luci/IR/AttrWeightCompression.h

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2024 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __LUCI_IR_ATTRWEIGHTCOMPRESSION_H__
+#define __LUCI_IR_ATTRWEIGHTCOMPRESSION_H__
+
+namespace luci
+{
+
+enum class WeightCompression
+{
+  UNDEFINED, // This is not defined by TFLite or Circle. This was added to
+             // prevent programming error.
+  NONE,
+  HUFFMAN
+};
+
+} // namespace luci
+
+#endif // __LUCI_IR_ATTRWEIGHTCOMPRESSION_H__

compiler/luci/lang/include/luci/IR/CircleNodeMixins.h

@@ -18,6 +18,7 @@
 #define __LUCI_IR_CIRCLE_NODE_MIXINS_H__
 
 #include "luci/IR/AttrFusedActFunc.h"
+#include "luci/IR/AttrWeightCompression.h"
 
 #include <loco/IR/Node.h>
 #include <loco/IR/NodeMixins.h>
@@ -31,11 +32,28 @@ namespace luci
 enum class CircleNodeTrait
 {
   FusedActFunc,
-  Bias
+  Bias,
+  WeightCompression
 };
 
 template <CircleNodeTrait T> class CircleNodeMixin;
 
+template <> class CircleNodeMixin<CircleNodeTrait::WeightCompression>
+{
+public:
+  CircleNodeMixin() = default;
+
+public:
+  WeightCompression weightCompression() const { return _weight_compression; }
+  void weightCompression(WeightCompression weight_compression)
+  {
+    _weight_compression = weight_compression;
+  }
+
+private:
+  WeightCompression _weight_compression = WeightCompression::UNDEFINED;
+};
+
 template <> class CircleNodeMixin<CircleNodeTrait::FusedActFunc>
 {
 public:

compiler/luci/lang/include/luci/IR/Nodes/CircleConv2D.h

@@ -34,7 +34,8 @@ namespace luci
  */
 class CircleConv2D final : public FixedArityNode<3, CircleNodeImpl<CircleOpcode::CONV_2D>>,
                            public CircleNodeMixin<CircleNodeTrait::FusedActFunc>,
-                           public CircleNodeMixin<CircleNodeTrait::Bias>
+                           public CircleNodeMixin<CircleNodeTrait::Bias>,
+                             public CircleNodeMixin<CircleNodeTrait::WeightCompression>
 {
 public:
   loco::Node *input(void) const { return at(0)->node(); }

compiler/luci/pass/include/luci/CircleOptimizer.h

@@ -109,6 +109,7 @@ class CircleOptimizer final
       UnrollUnidirSeqLSTM,
       XpSepActFromTransposeConv,
       RemoveGatherGuard,
+      CompressWeights
     };
 
     enum AlgorithmParameters

compiler/luci/pass/include/luci/Pass/CompressWeightsPass.h

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2024 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __LUCI_COMPRESS_WEIGHTS_PASS_H__
+#define __LUCI_COMPRESS_WEIGHTS_PASS_H__
+
+#include <logo/Pass.h>
+
+namespace luci
+{
+
+/**
+ * @brief  Class to generate FC/CONV with compressed weights
+ *
+ * To see the target Op pattern, please visit implementation.
+ */
+struct CompressWeightsPass final : public logo::Pass
+{
+  const char *name(void) const final { return "luci::CompressWeightsPass"; }
+
+  bool run(loco::Graph *g) final;
+};
+
+} // namespace luci
+
+#endif // __LUCI_COMPRESS_WEIGHTS_PASS_H__

compiler/luci/pass/src/CircleOptimizer.cpp

@@ -92,6 +92,7 @@
 #include "luci/Pass/DecomposeSoftmaxPass.h"
 #include "luci/Pass/UnrollUnidirectionalSequenceLSTMPass.h"
 #include "luci/Pass/XpSepActFromTransposeConvPass.h"
+#include "luci/Pass/CompressWeightsPass.h"
 // TODO add more passes
 
 #include "luci/Pass/CircleShapeInferencePass.h"
@@ -534,7 +535,10 @@ void CircleOptimizer::optimize(loco::Graph *g) const
   {
     phase.emplace_back(std::make_unique<luci::UnrollUnidirectionalSequenceLSTMPass>());
   }
-
+  if (_options->query(Options::Algorithm::CompressWeights))
+  {
+    phase.emplace_back(std::make_unique<luci::CompressWeightsPass>());
+  }
   // NOTE Experimental options; these will be removed someday
   //      Add experimental options here
   if (_options->query(Options::Algorithm::XpSepActFromTransposeConv))

compiler/luci/pass/src/CompressWeightsPass.cpp

@@ -0,0 +1,105 @@
+/*
+ * Copyright (c) 2024 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "luci/Pass/CompressWeightsPass.h"
+#include "helpers/HuffmanEncoder.h"
+#include "helpers/NodeFiller.h"
+
+#include <luci/IR/CircleNodes.h>
+#include <luci/Profile/CircleNodeOrigin.h>
+#include <luci/Service/CircleNodeClone.h>
+#include <luci/Service/Nodes/CircleConst.h>
+
+#include <cmath>
+#include <cassert>
+
+namespace
+{
+template <loco::DataType T> class TypeSelector;
+
+template <> class TypeSelector<loco::DataType::U8>
+{
+public:
+  using Type = uint8_t;
+};
+template <> class TypeSelector<loco::DataType::S8>
+{
+public:
+  using Type = int8_t;
+};
+
+template <loco::DataType DT> bool compress_weights_huffman_conv2d(luci::CircleConv2D *conv2d)
+{
+  using T = typename TypeSelector<DT>::Type;
+  assert(conv2d);
+  if (conv2d->weightCompression() == luci::WeightCompression::HUFFMAN)
+    return false;
+  luci::huffman::HuffmanEncoder<T> encoder;
+  auto weights = loco::must_cast<luci::CircleConst *>(conv2d->filter());
+  auto new_weights = luci::clone(weights);
+  std::vector<T> tmp_buf{};
+  for (int i = 0; i < weights->size<DT>(); ++i)
+  {
+    T data = weights->at<DT>(i);
+    tmp_buf.push_back(data);
+  }
+  std::vector<uint8_t> encoded = encoder.encode(tmp_buf);
+  new_weights->dtype(DT);
+  new_weights->size<DT>(encoded.size());
+  for (int i = 0; i < new_weights->size<DT>(); ++i)
+  {
+    new_weights->at<DT>(i) = encoded[i];
+  }
+  conv2d->filter(new_weights);
+  conv2d->weightCompression(luci::WeightCompression::HUFFMAN);
+
+  return true;
+}
+
+} // namespace
+
+namespace luci
+{
+
+bool CompressWeightsPass::run(loco::Graph *g)
+{
+  bool changed = false;
+  for (auto node : loco::active_nodes(loco::output_nodes(g)))
+  {
+    auto conv2d = dynamic_cast<luci::CircleConv2D *>(node);
+    if (not conv2d)
+      continue;
+    loco::DataType weights_dtype = loco::must_cast<luci::CircleConst *>(conv2d->filter())->dtype();
+    if (weights_dtype == loco::DataType::S8)
+    {
+      if (compress_weights_huffman_conv2d<loco::DataType::S8>(conv2d))
+        changed = true;
+    }
+    else if (weights_dtype == loco::DataType::U8)
+    {
+      if (compress_weights_huffman_conv2d<loco::DataType::U8>(conv2d))
+        changed = true;
+    }
+    else
+    {
+      throw std::runtime_error("Huffman weights compression supports s8 and u8");
+    }
+  }
+
+  return changed;
+}
+
+} // namespace luci