test: add LeNet MNIST CNN model e2e test (small)

PiperOrigin-RevId: 833907437

Author: asraa

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -154,6 +154,8 @@ void mlirToSecretArithmeticPipelineBuilder(
   pm.addPass(createActivationCanonicalizations());
   pm.addPass(createSelectRewrite());
   pm.addPass(createCompareToSignRewrite());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
 
   // Vectorize and optimize rotations
   // TODO(#2320): figure out where this fits in the new pipeline
@@ -493,8 +495,8 @@ BackendPipelineBuilder toLattigoPipelineBuilder() {
 
 void linalgPreprocessingBuilder(OpPassManager& manager) {
   manager.addPass(createInlineActivations());
-  manager.addPass(createDropUnitDims());
   manager.addPass(createLinalgCanonicalizations());
+  manager.addPass(createDropUnitDims());
   manager.addPass(createFoldConstantTensors());
   manager.addPass(createCanonicalizerPass());
   manager.addPass(createSymbolDCEPass());

lib/Target/OpenFhePke/Interpreter.cpp

@@ -175,6 +175,8 @@ std::vector<TypedCppValue> Interpreter::interpret(
 }
 
 void Interpreter::visit(Operation* op) {
+  std::cout << "Visiting operation: " << op->getName().getStringRef().str()
+            << "\n";
   llvm::TypeSwitch<Operation*>(op)
       .Case<arith::ConstantOp, arith::AddIOp, arith::SubIOp, arith::MulIOp,
             arith::MulFOp, arith::DivSIOp, arith::RemSIOp, arith::AndIOp,
@@ -198,6 +200,16 @@ void Interpreter::visit(Operation* op) {
         op->emitError() << "Unsupported operation " << opName.getStringRef()
                         << " in interpreter";
       });
+  // If any of the operations op operands have no more uses, then remove
+  // them from the end.
+  if (!op->getParentOfType<affine::AffineForOp>() &&
+      !op->getParentOfType<scf::ForOp>()) {
+    for (auto operand : op->getOperands()) {
+      if (liveness.isDeadAfter(operand, op)) {
+        env.erase(operand);
+      }
+    }
+  }
 }
 
 void Interpreter::visit(arith::ConstantOp op) {

lib/Transforms/LinalgCanonicalizations/LinalgCanonicalizations.cpp

@@ -517,9 +517,9 @@ struct RewriteAvgPoolAsConv2D
     RankedTensorType twoDInputType =
         RankedTensorType::get({inputTy.getDimSize(2), inputTy.getDimSize(3)},
                               inputTy.getElementType());
-    Value convOutput = rewriter.create<tensor::EmptyOp>(
-        poolOp.getLoc(), twoDOutputType.getShape(),
-        twoDOutputType.getElementType());
+    Value convOutput = tensor::EmptyOp::create(rewriter, poolOp.getLoc(),
+                                               twoDOutputType.getShape(),
+                                               twoDOutputType.getElementType());
     for (int n = 0; n < inputTy.getDimSize(0); ++n) {
       for (int c = 0; c < inputTy.getDimSize(1); ++c) {
         // Compute the 2-D constant convolution.
@@ -531,9 +531,9 @@ struct RewriteAvgPoolAsConv2D
             rewriter.getIndexAttr(inputTy.getDimSize(2)),
             rewriter.getIndexAttr(inputTy.getDimSize(3))};
         SmallVector<OpFoldResult> strides(4, rewriter.getIndexAttr(1));
-        auto extractInputOp = rewriter.create<tensor::ExtractSliceOp>(
-            poolOp.getLoc(), twoDInputType, poolOp.getInputs()[0], offsets,
-            inputSizes, strides);
+        auto extractInputOp = tensor::ExtractSliceOp::create(
+            rewriter, poolOp.getLoc(), twoDInputType, poolOp.getInputs()[0],
+            offsets, inputSizes, strides);
 
         auto convOp = linalg::Conv2DOp::create(
             rewriter, poolOp.getLoc(), twoDOutputType,
@@ -543,8 +543,8 @@ struct RewriteAvgPoolAsConv2D
             rewriter.getIndexAttr(1), rewriter.getIndexAttr(1),
             rewriter.getIndexAttr(outputTy.getDimSize(2)),
             rewriter.getIndexAttr(outputTy.getDimSize(3))};
-        outputVal = rewriter.create<tensor::InsertSliceOp>(
-            poolOp.getLoc(), convOp.getResult(0), outputVal, offsets,
+        outputVal = tensor::InsertSliceOp::create(
+            rewriter, poolOp.getLoc(), convOp.getResult(0), outputVal, offsets,
             outputSizes, strides);
       }
     }

lib/Utils/Layout/Codegen.cpp

@@ -212,13 +212,13 @@ Value buildIslExpr(isl_ast_expr* expr, std::map<std::string, Value> ivToValue,
         SmallVector<Value> args = getArgs(expr);
         auto op =
             arith::CmpIOp::create(b, islCmpToMlirAttr[type], args[0], args[1]);
-        return op->getResult(0);
-      }
-
-      if (type == isl_ast_op_select) {
+        auto indexCastOp = arith::IndexCastOp::create(b, b.getIndexType(), op);
+        return indexCastOp->getResult(0);
+      } else if (type == isl_ast_op_select) {
         // Select op
         SmallVector<Value> args = getArgs(expr);
-        auto op = arith::SelectOp::create(b, args[0], args[1], args[2]);
+        auto condI1 = arith::IndexCastOp::create(b, b.getI1Type(), args[0]);
+        auto op = arith::SelectOp::create(b, condI1, args[1], args[2]);
         return op->getResult(0);
       }
 
@@ -228,7 +228,9 @@ Value buildIslExpr(isl_ast_expr* expr, std::map<std::string, Value> ivToValue,
         auto op = arith::RemSIOp::create(b, args[0], args[1]);
         auto eqOp = arith::CmpIOp::create(b, arith::CmpIPredicate::eq, op,
                                           arith::ConstantIndexOp::create(b, 0));
-        return eqOp->getResult(0);
+        auto indexCastOp =
+            arith::IndexCastOp::create(b, b.getIndexType(), eqOp);
+        return indexCastOp->getResult(0);
       }
       isl_ast_expr_op_type opType = isl_ast_expr_get_op_type(expr);
       char* cStr = isl_ast_expr_to_C_str(expr);
@@ -317,9 +319,12 @@ FailureOr<scf::ValueVector> MLIRLoopNestGenerator::visitAstNodeIf(
   isl_ast_expr_free(cond);
 
   // Build scf if operation with the result types of the iter args
-  auto ifOp =
-      scf::IfOp::create(builder_, currentLoc_, TypeRange(currentIterArgs_),
-                        condVal, /*addThenBlock=*/true, /*addElseBlock=*/true);
+  // Convert condVal to an i1
+  auto condValI1 =
+      arith::IndexCastOp::create(builder_, builder_.getI1Type(), condVal);
+  auto ifOp = scf::IfOp::create(builder_, currentLoc_,
+                                TypeRange(currentIterArgs_), condValI1,
+                                /*addThenBlock=*/true, /*addElseBlock=*/true);
 
   // TODO:(#2120): Handle ISL else conditions.
   isl_ast_node* elseNode = isl_ast_node_if_get_else_node(node);

tests/Examples/openfhe/ckks/lenet/BUILD

@@ -0,0 +1,47 @@
+load("@heir//tests/Examples/openfhe:test.bzl", "openfhe_interpreter_test")
+
+# See README.md for setup required to run these tests
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+openfhe_interpreter_test(
+    name = "lenet_test",
+    timeout = "eternal",
+    generated_heir_opt_filename = "lenet_heir_opt.mlir",
+    heir_opt_flags = [
+        "--annotate-module=backend=openfhe scheme=ckks",
+        "--torch-linalg-to-ckks=ciphertext-degree=1024",
+        "--scheme-to-openfhe",
+    ],
+    mlir_src = "lenet.mlir",
+    tags = [
+        "manual",
+        "notap",
+    ],
+    test_src = "lenet_test.cpp",
+    deps = [
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+    ],
+)
+
+cc_binary(
+    name = "lenet_binary",
+    srcs = ["lenet_main.cpp"],
+    data = [
+        "lenet.openfhe.mlir",
+    ],
+    tags = [
+        "manual",
+        "notap",
+    ],
+    deps = [
+        "@bazel_tools//tools/cpp/runfiles",
+        "@heir//lib/Target/OpenFhePke:Interpreter",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@openfhe//:core",
+        "@openfhe//:pke",
+    ],
+)