Implement custom constant materializer to fix constant folding

larq_compute_engine/mlir/ir/lce_ops.cc

@@ -3,6 +3,7 @@
 #include "flatbuffers/flexbuffers.h"
 #include "larq_compute_engine/core/bitpacking/bitpack.h"
 #include "larq_compute_engine/mlir/transforms/bitpack.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
 // Generated dialect defs.
@@ -76,5 +77,12 @@ void LarqDialect::initialize() {
 #include "larq_compute_engine/mlir/ir/lce_ops.cc.inc"
       >();
 }
+
+Operation* LarqDialect::materializeConstant(OpBuilder& builder, Attribute value,
+                                            Type type, Location loc) {
+  if (arith::ConstantOp::isBuildableWith(value, type))
+    return builder.create<mlir::arith::ConstantOp>(loc, type, value);
+  return nullptr;
+}
 }  // namespace lq
 }  // namespace mlir

larq_compute_engine/mlir/ir/lce_ops.td

@@ -46,6 +46,8 @@ def LarqDialect : Dialect {
     TF graphs to be deployed on Larq Compute Engine.
   }];
 
+  let hasConstantMaterializer = 1;
+
   let cppNamespace = "::mlir::lq";
 }