Repository: google/heir
Pull request #1513: Plaintext backend: support mod arith semantic
Author: login: ZenithalHourlyRate
Labels: name: pull_ready
description: Indicates whether a PR is ready to pull. The copybara worker will import for internal testing
Created: 2025-03-04T03:29:48Z | Merged: 2025-03-12T16:16:23Z

After #1492 is in, we can now support arbitrary plaintext modulus execution in the plaintext backend, so we can know what happens when we using `i16` and `Zt`.

The notable example these two RUNs in `roberts_cross.mlir`, where bigger plaintext modulus will pass and smaller one will fail, indicating to the user/compiler that if the user/compiler chooses small plaintext modulus the failure is not from ciphertext side but from the Zt arithmetic.

Also this fixes #1509 by using `cc` from environment.

Linked issue #1509: tests: should find clang from env
After #1466 is in, clang now becomes a dependency when running `bazel test //tests/Examples/cleartext/...`.

We should make lit access `clang` from environment.
Linked issue comments: user: login: j2kun
created_at: 2025-03-03T21:14:51Z
body: I wonder if it would be possible to replace `clang` in these tests with `$CC` or `$CXX` and let it query the (hopefully inherited) environment variables.
author_association: COLLABORATOR
user: login: ZenithalHourlyRate
created_at: 2025-03-04T01:31:13Z
body: ```
# RUN: at line 3
bash -c 'echo $PATH'
# executed command: bash -c 'echo $PATH'
# .---command stdout------------
# | /run/user/2015/T/bazel/88da44a58ff2b8ea788068322eb5e685/sandbox/linux-sandbox/5035/execroot/_main/bazel-out/k8-dbg/bin/tests/Examples/cleartext/dot_product_8.mlir.test.runfiles/llvm-project/mlir:/run/user/2015/T/bazel/88da44a58ff2b8ea788068322eb5e685/sandbox/linux-sandbox/5035/execroot/_main/bazel-out/k8-dbg/bin/tests/Examples/cleartext/dot_product_8.mlir.test.runfiles/_main/tools:/run/user/2015/T/bazel/88da44a58ff2b8ea788068322eb5e685/sandbox/linux-sandbox/5035/execroot/_main/bazel-out/k8-dbg/bin/tests/Examples/cleartext/dot_product_8.mlir.test.runfiles/_main/tests/Emitter/verilog:/run/user/2015/T/bazel/88da44a58ff2b8ea788068322eb5e685/sandbox/linux-sandbox/5035/execroot/_main/bazel-out/k8-dbg/bin/tests/Examples/cleartext/dot_product_8.mlir.test.runfiles/llvm-project/llvm:/run/user/2015/T/bazel/88da44a58ff2b8ea788068322eb5e685/sandbox/linux-sandbox/5035/execroot/_main/bazel-out/k8-dbg/bin/tests/Examples/cleartext/dot_product_8.mlir.test.runfiles/at_clifford_yosys:.:/bin:/usr/bin:/usr/local/bin
```

Seems that within the lit test we can access things in `/bin`. Other environment variable are not inherited unless we change `lit.cfg.py` pass that variable and does detection logic of which `$CC` should we use (user usually does not set that, and lit may could query bazel for it though)

I think default to `/usr/bin/cc` and `/usr/bin/c++` would be fair enough as any machine that could run `bazel test` should have that.
author_association: COLLABORATOR

Diff:
diff --git a/lib/Pipelines/ArithmeticPipelineRegistration.cpp b/lib/Pipelines/ArithmeticPipelineRegistration.cpp
index 798937a0f3..e071cd67c3 100644
--- a/lib/Pipelines/ArithmeticPipelineRegistration.cpp
+++ b/lib/Pipelines/ArithmeticPipelineRegistration.cpp
@@ -4,6 +4,7 @@
 #include <string>
 #include <vector>
 
+#include "lib/Dialect/Arith/Conversions/ArithToModArith/ArithToModArith.h"
 #include "lib/Dialect/BGV/Conversions/BGVToLWE/BGVToLWE.h"
 #include "lib/Dialect/CKKS/Conversions/CKKSToLWE/CKKSToLWE.h"
 #include "lib/Dialect/LWE/Conversions/LWEToLattigo/LWEToLattigo.h"
@@ -18,7 +19,9 @@
 #include "lib/Dialect/Secret/Conversions/SecretToBGV/SecretToBGV.h"
 #include "lib/Dialect/Secret/Conversions/SecretToCKKS/SecretToCKKS.h"
 #include "lib/Dialect/Secret/Transforms/DistributeGeneric.h"
+#include "lib/Dialect/Secret/Transforms/ForgetSecrets.h"
 #include "lib/Dialect/Secret/Transforms/MergeAdjacentGenerics.h"
+#include "lib/Dialect/TensorExt/Conversions/TensorExtToTensor/TensorExtToTensor.h"
 #include "lib/Dialect/TensorExt/Transforms/CollapseInsertionChains.h"
 #include "lib/Dialect/TensorExt/Transforms/InsertRotate.h"
 #include "lib/Dialect/TensorExt/Transforms/RotateAndReduce.h"
@@ -123,6 +126,35 @@ void mlirToSecretArithmeticPipelineBuilder(
   pm.addPass(createOperationBalancer());
 }
 
+void mlirToPlaintextPipelineBuilder(OpPassManager &pm,
+                                    const PlaintextBackendOptions &options) {
+  // Convert to secret arithmetic
+  MlirToRLWEPipelineOptions mlirToRLWEPipelineOptions;
+  mlirToSecretArithmeticPipelineBuilder(pm, mlirToRLWEPipelineOptions);
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+
+  // Forget secrets to convert secret types to standard types
+  pm.addPass(secret::createSecretForgetSecrets());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+
+  if (options.plaintextModulus != 0) {
+    // Convert to plaintext arithmetic in Zp
+    arith::ArithToModArithOptions arithToModArithOptions;
+    arithToModArithOptions.modulus = options.plaintextModulus;
+    pm.addPass(arith::createArithToModArith(arithToModArithOptions));
+    pm.addPass(createCanonicalizerPass());
+    pm.addPass(createCSEPass());
+  }
+
+  // Convert to standard dialect
+  pm.addPass(tensor_ext::createTensorExtToTensor());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+  polynomialToLLVMPipelineBuilder(pm);
+}
+
 void mlirToRLWEPipeline(OpPassManager &pm,
                         const MlirToRLWEPipelineOptions &options,
                         const RLWEScheme scheme) {
diff --git a/lib/Pipelines/ArithmeticPipelineRegistration.h b/lib/Pipelines/ArithmeticPipelineRegistration.h
index 7d3b1a95d0..c5393a1439 100644
--- a/lib/Pipelines/ArithmeticPipelineRegistration.h
+++ b/lib/Pipelines/ArithmeticPipelineRegistration.h
@@ -68,6 +68,15 @@ struct MlirToRLWEPipelineOptions : public SimdVectorizerOptions {
       llvm::cl::init(45)};
 };
 
+struct PlaintextBackendOptions
+    : public PassPipelineOptions<PlaintextBackendOptions> {
+  PassOptions::Option<int64_t> plaintextModulus{
+      *this, "plaintext-modulus",
+      llvm::cl::desc("Plaintext modulus for BGV/BFV scheme (if not specified, "
+                     "execute in the original integer type)"),
+      llvm::cl::init(0)};
+};
+
 struct BackendOptions : public PassPipelineOptions<BackendOptions> {
   PassOptions::Option<std::string> entryFunction{
       *this, "entry-function", llvm::cl::desc("Entry function"),
@@ -92,6 +101,9 @@ void mlirToRLWEPipeline(OpPassManager &pm,
 void mlirToSecretArithmeticPipelineBuilder(
     OpPassManager &pm, const MlirToRLWEPipelineOptions &options);
 
+void mlirToPlaintextPipelineBuilder(OpPassManager &pm,
+                                    const PlaintextBackendOptions &options);
+
 RLWEPipelineBuilder mlirToRLWEPipelineBuilder(RLWEScheme scheme);
 
 BackendPipelineBuilder toOpenFhePipelineBuilder();
diff --git a/lib/Pipelines/BUILD b/lib/Pipelines/BUILD
index 59c2e427c8..013d58c09b 100644
--- a/lib/Pipelines/BUILD
+++ b/lib/Pipelines/BUILD
@@ -85,6 +85,7 @@ cc_library(
     hdrs = ["ArithmeticPipelineRegistration.h"],
     deps = [
         ":PipelineRegistration",
+        "@heir//lib/Dialect/Arith/Conversions/ArithToModArith",
         "@heir//lib/Dialect/BGV/Conversions/BGVToLWE",
         "@heir//lib/Dialect/CKKS/Conversions/CKKSToLWE",
         "@heir//lib/Dialect/LWE/Conversions/LWEToLattigo",
@@ -101,8 +102,10 @@ cc_library(
         "@heir//lib/Dialect/Secret/Conversions/SecretToCGGI",
         "@heir//lib/Dialect/Secret/Conversions/SecretToCKKS",
         "@heir//lib/Dialect/Secret/Transforms:DistributeGeneric",
+        "@heir//lib/Dialect/Secret/Transforms:ForgetSecrets",
         "@heir//lib/Dialect/Secret/Transforms:MergeAdjacentGenerics",
         "@heir//lib/Dialect/TOSA/Conversions/TosaToSecretArith",
+        "@heir//lib/Dialect/TensorExt/Conversions/TensorExtToTensor",
         "@heir//lib/Dialect/TensorExt/Transforms:CollapseInsertionChains",
         "@heir//lib/Dialect/TensorExt/Transforms:InsertRotate",
         "@heir//lib/Dialect/TensorExt/Transforms:RotateAndReduce",
diff --git a/lib/Pipelines/PipelineRegistration.cpp b/lib/Pipelines/PipelineRegistration.cpp
index aa29104992..e4670e7116 100644
--- a/lib/Pipelines/PipelineRegistration.cpp
+++ b/lib/Pipelines/PipelineRegistration.cpp
@@ -126,6 +126,11 @@ void polynomialToLLVMPipelineBuilder(OpPassManager &manager) {
   manager.addPass(arith::createArithExpandOpsPass());
   manager.addPass(createSCFToControlFlowPass());
   manager.addNestedPass<FuncOp>(memref::createExpandStridedMetadataPass());
+  // expand strided metadata will create affine map. Needed to lower affine.map
+  // and affine.apply
+  manager.addNestedPass<FuncOp>(affine::createAffineExpandIndexOpsPass());
+  manager.addNestedPass<FuncOp>(affine::createSimplifyAffineStructuresPass());
+  manager.addPass(createLowerAffinePass());
   manager.addPass(createConvertToLLVMPass());
 
   // Cleanup
diff --git a/tests/BUILD b/tests/BUILD
index 96647fc1fc..b6fea66f21 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -36,7 +36,7 @@ glob_lit_tests(
     exclude = [
         "Examples/common/*",
         # Has a custom driver
-        "Examples/cleartext/dot_product_8.mlir",
+        "Examples/plaintext/*",
     ],
     test_file_exts = ["mlir"],
 )
diff --git a/tests/Examples/cleartext/BUILD b/tests/Examples/cleartext/BUILD
deleted file mode 100644
index 6e2f03cb13..0000000000
--- a/tests/Examples/cleartext/BUILD
+++ /dev/null
@@ -1,24 +0,0 @@
-load("//bazel:lit.bzl", "glob_lit_tests")
-
-package(default_applicable_licenses = ["@heir//:license"])
-
-glob_lit_tests(
-    name = "all_tests",
-    data = [
-        ":dot_product_8_test",
-        "@heir//tests:test_utilities",
-        "@llvm-project//clang",
-        "@llvm-project//llvm:llc",
-        "@llvm-project//mlir:mlir-translate",
-    ],
-    default_tags = [
-        "notap",
-    ],
-    driver = "@heir//tests:run_lit.sh",
-    test_file_exts = ["mlir"],
-)
-
-cc_library(
-    name = "dot_product_8_test",
-    srcs = ["dot_product_8_test.c"],
-)
diff --git a/tests/Examples/cleartext/dot_product_8.mlir b/tests/Examples/cleartext/dot_product_8.mlir
deleted file mode 100644
index 28b51db494..0000000000
--- a/tests/Examples/cleartext/dot_product_8.mlir
+++ /dev/null
@@ -1,18 +0,0 @@
-// RUN: heir-opt %s --mlir-to-secret-arithmetic --tensor-ext-to-tensor --secret-forget-secrets --heir-polynomial-to-llvm \
-// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
-// RUN: clang %S/libdot_product_8_test.a %t -o a.out
-// RUN: ./a.out | FileCheck %s
-
-// CHECK: Test passed
-func.func @dot_product(%arg0: tensor<8xi16> {secret.secret}, %arg1: tensor<8xi16> {secret.secret}) -> i16 {
-  %c0 = arith.constant 0 : index
-  %c0_si16 = arith.constant 0 : i16
-  %0 = affine.for %arg2 = 0 to 8 iter_args(%iter = %c0_si16) -> (i16) {
-    %1 = tensor.extract %arg0[%arg2] : tensor<8xi16>
-    %2 = tensor.extract %arg1[%arg2] : tensor<8xi16>
-    %3 = arith.muli %1, %2 : i16
-    %4 = arith.addi %iter, %3 : i16
-    affine.yield %4 : i16
-  }
-  return %0 : i16
-}
diff --git a/tests/Examples/plaintext/BUILD b/tests/Examples/plaintext/BUILD
new file mode 100644
index 0000000000..78658a21fc
--- /dev/null
+++ b/tests/Examples/plaintext/BUILD
@@ -0,0 +1 @@
+package(default_applicable_licenses = ["@heir//:license"])
diff --git a/tests/Examples/plaintext/dot_product/BUILD b/tests/Examples/plaintext/dot_product/BUILD
new file mode 100644
index 0000000000..5bbfc8c4f2
--- /dev/null
+++ b/tests/Examples/plaintext/dot_product/BUILD
@@ -0,0 +1,21 @@
+load("//tests/Examples/plaintext:test.bzl", "plaintext_test")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+plaintext_test(
+    src = "dot_product_8.mlir",
+    data = [
+        ":dot_product_8_mod_test",
+        ":dot_product_8_test",
+    ],
+)
+
+cc_library(
+    name = "dot_product_8_test",
+    srcs = ["dot_product_8_test.c"],
+)
+
+cc_library(
+    name = "dot_product_8_mod_test",
+    srcs = ["dot_product_8_mod_test.c"],
+)
diff --git a/tests/Examples/plaintext/dot_product/dot_product_8.mlir b/tests/Examples/plaintext/dot_product/dot_product_8.mlir
new file mode 100644
index 0000000000..d1dd978ba6
--- /dev/null
+++ b/tests/Examples/plaintext/dot_product/dot_product_8.mlir
@@ -0,0 +1,42 @@
+// Secret arithmetic plaintext backend with original integer type semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_test.a %t -o a.out
+// RUN: ./a.out | FileCheck %s
+
+// Secret arithmetic plaintext backend with plaintext modulus semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend=plaintext-modulus=257 \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_mod_test.a %t -o a.out
+// RUN: ./a.out | FileCheck %s
+
+// Secret arithmetic plaintext backend with plaintext modulus semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend=plaintext-modulus=65537 \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_mod_test.a %t -o a.out
+// RUN: ./a.out | FileCheck %s
+
+// Secret arithmetic plaintext backend with plaintext modulus semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend=plaintext-modulus=17 \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_mod_test.a %t -o a.out
+// RUN: ./a.out | FileCheck %s --check-prefix=CHECK-SMALL
+
+// CHECK: Test passed
+// CHECK-SMALL: Test failed 2 != 240
+func.func @dot_product(%arg0: tensor<8xi16> {secret.secret}, %arg1: tensor<8xi16> {secret.secret}) -> i16 {
+  %c0 = arith.constant 0 : index
+  %c0_si16 = arith.constant 0 : i16
+  %0 = affine.for %arg2 = 0 to 8 iter_args(%iter = %c0_si16) -> (i16) {
+    %1 = tensor.extract %arg0[%arg2] : tensor<8xi16>
+    %2 = tensor.extract %arg1[%arg2] : tensor<8xi16>
+    %3 = arith.muli %1, %2 : i16
+    %4 = arith.addi %iter, %3 : i16
+    affine.yield %4 : i16
+  }
+  return %0 : i16
+}
diff --git a/tests/Examples/plaintext/dot_product/dot_product_8_mod_test.c b/tests/Examples/plaintext/dot_product/dot_product_8_mod_test.c
new file mode 100644
index 0000000000..47b2278629
--- /dev/null
+++ b/tests/Examples/plaintext/dot_product/dot_product_8_mod_test.c
@@ -0,0 +1,30 @@
+#include <stdint.h>
+#include <stdio.h>
+
+// This is the function we want to call from LLVM
+// For mod-arith, the width is 64 bits
+int64_t dot_product(
+    /* arg 0*/
+    int64_t *allocated, int64_t *aligned, int64_t offset, int64_t size,
+    int64_t stride,
+    /* arg 1*/
+    int64_t *allocated_2, int64_t *aligned_2, int64_t offset_2, int64_t size_2,
+    int64_t stride_2);
+
+int main() {
+  int64_t arg0[8] = {1, 2, 3, 4, 5, 6, 7, 8};
+  int64_t arg1[8] = {2, 3, 4, 5, 6, 7, 8, 9};
+  int64_t expected = 240;
+
+  int64_t res = dot_product(
+      /* arg 0*/
+      arg0, arg0, 0, 8, 1,
+      /* arg 1*/
+      arg1, arg1, 0, 8, 1);
+
+  if (res == expected) {
+    printf("Test passed\n");
+  } else {
+    printf("Test failed %ld != %ld\n", res, expected);
+  }
+}
diff --git a/tests/Examples/cleartext/dot_product_8_test.c b/tests/Examples/plaintext/dot_product/dot_product_8_test.c
similarity index 100%
rename from tests/Examples/cleartext/dot_product_8_test.c
rename to tests/Examples/plaintext/dot_product/dot_product_8_test.c
diff --git a/tests/Examples/plaintext/dot_product_f/BUILD b/tests/Examples/plaintext/dot_product_f/BUILD
new file mode 100644
index 0000000000..6b2942e906
--- /dev/null
+++ b/tests/Examples/plaintext/dot_product_f/BUILD
@@ -0,0 +1,15 @@
+load("//tests/Examples/plaintext:test.bzl", "plaintext_test")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+plaintext_test(
+    src = "dot_product_8f.mlir",
+    data = [
+        ":dot_product_8f_test",
+    ],
+)
+
+cc_library(
+    name = "dot_product_8f_test",
+    srcs = ["dot_product_8f_test.c"],
+)
diff --git a/tests/Examples/plaintext/dot_product_f/dot_product_8f.mlir b/tests/Examples/plaintext/dot_product_f/dot_product_8f.mlir
new file mode 100644
index 0000000000..8ab38eb607
--- /dev/null
+++ b/tests/Examples/plaintext/dot_product_f/dot_product_8f.mlir
@@ -0,0 +1,20 @@
+// Secret arithmetic plaintext backend with original integer type semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8f_test.a %t -o a.out
+// RUN: ./a.out | FileCheck %s
+
+// CHECK: Test passed
+func.func @dot_product(%arg0: tensor<8xf32> {secret.secret}, %arg1: tensor<8xf32> {secret.secret}) -> f32 {
+  %c0 = arith.constant 0 : index
+  %c0_sf32 = arith.constant 0.1 : f32
+  %0 = affine.for %arg2 = 0 to 8 iter_args(%iter = %c0_sf32) -> (f32) {
+    %1 = tensor.extract %arg0[%arg2] : tensor<8xf32>
+    %2 = tensor.extract %arg1[%arg2] : tensor<8xf32>
+    %3 = arith.mulf %1, %2 : f32
+    %4 = arith.addf %iter, %3 : f32
+    affine.yield %4 : f32
+  }
+  return %0 : f32
+}
diff --git a/tests/Examples/plaintext/dot_product_f/dot_product_8f_test.c b/tests/Examples/plaintext/dot_product_f/dot_product_8f_test.c
new file mode 100644
index 0000000000..37f4d71448
--- /dev/null
+++ b/tests/Examples/plaintext/dot_product_f/dot_product_8f_test.c
@@ -0,0 +1,30 @@
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+
+// This is the function we want to call from LLVM
+float dot_product(
+    /* arg 0*/
+    float *allocated, float *aligned, int64_t offset, int64_t size,
+    int64_t stride,
+    /* arg 1*/
+    float *allocated_2, float *aligned_2, int64_t offset_2, int64_t size_2,
+    int64_t stride_2);
+
+int main() {
+  float arg0[8] = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8};
+  float arg1[8] = {0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9};
+  float expected = 2.5;
+
+  float res = dot_product(
+      /* arg 0*/
+      arg0, arg0, 0, 8, 1,
+      /* arg 1*/
+      arg1, arg1, 0, 8, 1);
+
+  if (fabs(res - expected) < 1e-3) {
+    printf("Test passed\n");
+  } else {
+    printf("Test failed %f != %f\n", res, expected);
+  }
+}
diff --git a/tests/Examples/plaintext/roberts_cross/BUILD b/tests/Examples/plaintext/roberts_cross/BUILD
new file mode 100644
index 0000000000..de7aa44ed1
--- /dev/null
+++ b/tests/Examples/plaintext/roberts_cross/BUILD
@@ -0,0 +1,33 @@
+load("//tests/Examples/plaintext:test.bzl", "plaintext_test")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+plaintext_test(
+    size = "enormous",
+    src = "roberts_cross_64x64.mlir",
+    data = [
+        ":roberts_cross_mod_test",
+        ":roberts_cross_test",
+    ],
+    tags = [
+        "manual",
+        "nofastbuild",
+        "notap",
+    ],
+)
+
+cc_library(
+    name = "roberts_cross_test",
+    srcs = [
+        "memrefCopy.cpp",
+        "roberts_cross_test.c",
+    ],
+)
+
+cc_library(
+    name = "roberts_cross_mod_test",
+    srcs = [
+        "memrefCopy.cpp",
+        "roberts_cross_mod_test.c",
+    ],
+)
diff --git a/tests/Examples/plaintext/roberts_cross/memrefCopy.cpp b/tests/Examples/plaintext/roberts_cross/memrefCopy.cpp
new file mode 100644
index 0000000000..d4b00d5a57
--- /dev/null
+++ b/tests/Examples/plaintext/roberts_cross/memrefCopy.cpp
@@ -0,0 +1,100 @@
+#include <cstdlib>
+#include <cstring>
+
+// Copied from CRunnerUtils from MLIR
+
+/// StridedMemRef descriptor type with static rank.
+template <typename T, int N>
+struct StridedMemRefType {
+  T *basePtr;
+  T *data;
+  int64_t offset;
+  int64_t sizes[N];
+  int64_t strides[N];
+};
+
+//===----------------------------------------------------------------------===//
+// Codegen-compatible structure for UnrankedMemRef type.
+//===----------------------------------------------------------------------===//
+// Unranked MemRef
+template <typename T>
+struct UnrankedMemRefType {
+  int64_t rank;
+  void *descriptor;
+};
+
+// A reference to one of the StridedMemRef types.
+template <typename T>
+class DynamicMemRefType {
+ public:
+  int64_t rank;
+  T *basePtr;
+  T *data;
+  int64_t offset;
+  const int64_t *sizes;
+  const int64_t *strides;
+
+  explicit DynamicMemRefType(const ::UnrankedMemRefType<T> &memRef)
+      : rank(memRef.rank) {
+    auto *desc = static_cast<StridedMemRefType<T, 1> *>(memRef.descriptor);
+    basePtr = desc->basePtr;
+    data = desc->data;
+    offset = desc->offset;
+    sizes = rank == 0 ? nullptr : desc->sizes;
+    strides = sizes + rank;
+  }
+};
+
+extern "C" void memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
+                           UnrankedMemRefType<char> *dstArg) {
+  DynamicMemRefType<char> src(*srcArg);
+  DynamicMemRefType<char> dst(*dstArg);
+
+  int64_t rank = src.rank;
+  // MLIR_MSAN_MEMORY_IS_INITIALIZED(src.sizes, rank * sizeof(int64_t));
+
+  // Handle empty shapes -> nothing to copy.
+  for (int rankp = 0; rankp < rank; ++rankp)
+    if (src.sizes[rankp] == 0) return;
+
+  char *srcPtr = src.data + src.offset * elemSize;
+  char *dstPtr = dst.data + dst.offset * elemSize;
+
+  if (rank == 0) {
+    memcpy(dstPtr, srcPtr, elemSize);
+    return;
+  }
+
+  int64_t *indices = static_cast<int64_t *>(alloca(sizeof(int64_t) * rank));
+  int64_t *srcStrides = static_cast<int64_t *>(alloca(sizeof(int64_t) * rank));
+  int64_t *dstStrides = static_cast<int64_t *>(alloca(sizeof(int64_t) * rank));
+
+  // Initialize index and scale strides.
+  for (int rankp = 0; rankp < rank; ++rankp) {
+    indices[rankp] = 0;
+    srcStrides[rankp] = src.strides[rankp] * elemSize;
+    dstStrides[rankp] = dst.strides[rankp] * elemSize;
+  }
+
+  int64_t readIndex = 0, writeIndex = 0;
+  for (;;) {
+    // Copy over the element, byte by byte.
+    memcpy(dstPtr + writeIndex, srcPtr + readIndex, elemSize);
+    // Advance index and read position.
+    for (int64_t axis = rank - 1; axis >= 0; --axis) {
+      // Advance at current axis.
+      auto newIndex = ++indices[axis];
+      readIndex += srcStrides[axis];
+      writeIndex += dstStrides[axis];
+      // If this is a valid index, we have our next index, so continue copying.
+      if (src.sizes[axis] != newIndex) break;
+      // We reached the end of this axis. If this is axis 0, we are done.
+      if (axis == 0) return;
+      // Else, reset to 0 and undo the advancement of the linear index that
+      // this axis had. Then continue with the axis one outer.
+      indices[axis] = 0;
+      readIndex -= src.sizes[axis] * srcStrides[axis];
+      writeIndex -= dst.sizes[axis] * dstStrides[axis];
+    }
+  }
+}
diff --git a/tests/Examples/plaintext/roberts_cross/roberts_cross_64x64.mlir b/tests/Examples/plaintext/roberts_cross/roberts_cross_64x64.mlir
new file mode 100644
index 0000000000..835f70186c
--- /dev/null
+++ b/tests/Examples/plaintext/roberts_cross/roberts_cross_64x64.mlir
@@ -0,0 +1,84 @@
+// Secret arithmetic plaintext backend with original integer type semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %t %S/libroberts_cross_test.a -o a.out
+// RUN: ./a.out | FileCheck %s --check-prefix=CHECK-INT
+// CHECK-INT: Test passed
+
+// Secret arithmetic plaintext backend with plaintext modulus semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend=plaintext-modulus=536903681 \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %t %S/libroberts_cross_mod_test.a -o a.out
+// RUN: ./a.out | FileCheck %s
+// CHECK: Test passed
+
+// Secret arithmetic plaintext backend with plaintext modulus semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend=plaintext-modulus=786433 \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %t %S/libroberts_cross_mod_test.a -o a.out
+// RUN: ./a.out | FileCheck %s --check-prefix=CHECK-SMALL
+// CHECK-SMALL: Test failed
+func.func @roberts_cross(%img: tensor<4096xi16> {secret.secret}) -> tensor<4096xi16> {
+  %c4096 = arith.constant 4096 : index
+  %c64 = arith.constant 64 : index
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %c-1 =  arith.constant -1 : index
+
+  // Each point p = img[x][y], where x is row and y is column, in the new image will equal:
+  // (img[x-1][y-1] - img[x][y])^2 + (img[x-1][y] - img[x][y-1])^2
+  %r = affine.for %x = 0 to 64 iter_args(%imgx = %img) -> tensor<4096xi16> {
+    %1 = affine.for %y = 0 to 64 iter_args(%imgy = %imgx) -> tensor<4096xi16> {
+
+      // fetch img[x-1][y-1]
+      %4 = arith.addi %x, %c-1 : index
+      %5 = arith.muli %4, %c64 : index
+      %6 = arith.addi %y, %c-1 : index
+      %7 = arith.addi %5, %6 : index
+      %8 = arith.remui %7, %c4096 : index
+      %9 = tensor.extract %img[%8] : tensor<4096xi16>
+
+      // fetch img[x][y]
+      %10 = arith.muli %x, %c64 : index
+      %11 = arith.addi %10, %y : index
+      %12 = arith.remui %11, %c4096 : index
+      %13 = tensor.extract %img[%12] : tensor<4096xi16>
+
+      // subtract those two
+      %14 = arith.subi %9, %13 : i16
+
+      // fetch img[x-1][y]
+      %15 = arith.addi %x, %c-1 : index
+      %16 = arith.muli %15, %c64 : index
+      %18 = arith.addi %16, %y : index
+      %19 = arith.remui %18, %c4096 : index
+      %20 = tensor.extract %img[%19] : tensor<4096xi16>
+
+      // fetch img[x][y-1]
+      %21 = arith.muli %x, %c64 : index
+      %22 = arith.addi %y, %c-1 : index
+      %23 = arith.addi %21, %22 : index
+      %24 = arith.remui %23, %c4096 : index
+      %25 = tensor.extract %img[%24] : tensor<4096xi16>
+
+      // subtract those two
+      %26 = arith.subi %20, %25 : i16
+
+      // square each difference
+      %27 = arith.muli %14, %14 :  i16
+      %28 = arith.muli %26, %26 :  i16
+
+      // add the squares
+      %29 = arith.addi %27, %28 : i16
+
+      // save to result[x][y]
+      %30 = tensor.insert %29 into %imgy[%12] : tensor<4096xi16>
+      affine.yield %30: tensor<4096xi16>
+    }
+    affine.yield %1 : tensor<4096xi16>
+  }
+  return %r : tensor<4096xi16>
+}
diff --git a/tests/Examples/plaintext/roberts_cross/roberts_cross_mod_test.c b/tests/Examples/plaintext/roberts_cross/roberts_cross_mod_test.c
new file mode 100644
index 0000000000..9c42768b8a
--- /dev/null
+++ b/tests/Examples/plaintext/roberts_cross/roberts_cross_mod_test.c
@@ -0,0 +1,61 @@
+#include <stdint.h>
+#include <stdio.h>
+
+struct Memref1D {
+  int64_t *allocated;
+  int64_t *aligned;
+  int64_t offset;
+  int64_t size;
+  int64_t stride;
+};
+
+// This is the function we want to call from LLVM
+// For mod-arith, the width is 64 bits
+struct Memref1D roberts_cross(
+    /* arg0*/
+    int64_t *allocated, int64_t *aligned, int64_t offset, int64_t size,
+    int64_t stride);
+
+void memrefCopy();
+
+int main() {
+  int64_t input[4096];
+  int64_t expected[4096];
+
+  for (int i = 0; i < 4096; ++i) {
+    input[i] = i;
+  }
+
+  for (int row = 0; row < 64; ++row) {
+    for (int col = 0; col < 64; ++col) {
+      // (img[x-1][y-1] - img[x][y])^2 + (img[x-1][y] - img[x][y-1])^2
+      int64_t xY = (row * 64 + col) % 4096;
+      int64_t xYm1 = (row * 64 + col - 1) % 4096;
+      int64_t xm1Y = ((row - 1) * 64 + col) % 4096;
+      int64_t xm1Ym1 = ((row - 1) * 64 + col - 1) % 4096;
+
+      if (xYm1 < 0) xYm1 += 4096;
+      if (xm1Y < 0) xm1Y += 4096;
+      if (xm1Ym1 < 0) xm1Ym1 += 4096;
+
+      int64_t v1 = (input[xm1Ym1] - input[xY]);
+      int64_t v2 = (input[xm1Y] - input[xYm1]);
+      int64_t sum = v1 * v1 + v2 * v2;
+      expected[row * 64 + col] = sum;
+    }
+  }
+
+  struct Memref1D memref = roberts_cross(
+      /* arg 0*/
+      input, input, 0, 4096, 1);
+
+  int64_t *res = memref.aligned;
+
+  for (int i = 0; i != 4096; ++i) {
+    if (res[i] != expected[i]) {
+      printf("Test failed at %d: %ld != %ld\n", i, res[i], expected[i]);
+      return 0;
+    }
+  }
+  printf("Test passed\n");
+}
diff --git a/tests/Examples/plaintext/roberts_cross/roberts_cross_test.c b/tests/Examples/plaintext/roberts_cross/roberts_cross_test.c
new file mode 100644
index 0000000000..68a13007e1
--- /dev/null
+++ b/tests/Examples/plaintext/roberts_cross/roberts_cross_test.c
@@ -0,0 +1,61 @@
+#include <stdint.h>
+#include <stdio.h>
+
+struct Memref1D {
+  int16_t *allocated;
+  int16_t *aligned;
+  int64_t offset;
+  int64_t size;
+  int64_t stride;
+};
+
+// This is the function we want to call from LLVM
+// For mod-arith, the width is 64 bits
+struct Memref1D roberts_cross(
+    /* arg0*/
+    int16_t *allocated, int16_t *aligned, int64_t offset, int64_t size,
+    int64_t stride);
+
+void memrefCopy();
+
+int main() {
+  int16_t input[4096];
+  int16_t expected[4096];
+
+  for (int i = 0; i < 4096; ++i) {
+    input[i] = i;
+  }
+
+  for (int row = 0; row < 64; ++row) {
+    for (int col = 0; col < 64; ++col) {
+      // (img[x-1][y-1] - img[x][y])^2 + (img[x-1][y] - img[x][y-1])^2
+      int16_t xY = (row * 64 + col) % 4096;
+      int16_t xYm1 = (row * 64 + col - 1) % 4096;
+      int16_t xm1Y = ((row - 1) * 64 + col) % 4096;
+      int16_t xm1Ym1 = ((row - 1) * 64 + col - 1) % 4096;
+
+      if (xYm1 < 0) xYm1 += 4096;
+      if (xm1Y < 0) xm1Y += 4096;
+      if (xm1Ym1 < 0) xm1Ym1 += 4096;
+
+      int16_t v1 = (input[xm1Ym1] - input[xY]);
+      int16_t v2 = (input[xm1Y] - input[xYm1]);
+      int16_t sum = v1 * v1 + v2 * v2;
+      expected[row * 64 + col] = sum;
+    }
+  }
+
+  struct Memref1D memref = roberts_cross(
+      /* arg 0*/
+      input, input, 0, 4096, 1);
+
+  int16_t *res = memref.aligned;
+
+  for (int i = 0; i != 4096; ++i) {
+    if (res[i] != expected[i]) {
+      printf("Test failed at %d: %d != %d\n", i, res[i], expected[i]);
+      return 0;
+    }
+  }
+  printf("Test passed\n");
+}
diff --git a/tests/Examples/plaintext/test.bzl b/tests/Examples/plaintext/test.bzl
new file mode 100644
index 0000000000..57d0abb9bb
--- /dev/null
+++ b/tests/Examples/plaintext/test.bzl
@@ -0,0 +1,30 @@
+"""A macro providing an end-to-end test for Plaintext Backend."""
+
+load("@heir//bazel:lit.bzl", "lit_test")
+
+def plaintext_test(src, size = "small", data = [], tags = []):
+    """Define a lit test for the Plaintext Backend.
+
+    Args:
+      src: The source mlir file to run through lit.
+      size: The size of the test.
+      data: Data dependencies to be passed to lit_test
+      tags: Tags to pass to lit_test
+    """
+
+    _data = data + [
+        "@heir//tests:test_utilities",
+        "@llvm-project//llvm:llc",
+        "@llvm-project//mlir:mlir-translate",
+    ]
+
+    _tags = tags + [
+        "notap",
+    ]
+
+    lit_test(
+        src = src,
+        size = size,
+        data = _data,
+        tags = _tags,
+    )
diff --git a/tools/heir-opt.cpp b/tools/heir-opt.cpp
index 771ba2bddb..2565676b43 100644
--- a/tools/heir-opt.cpp
+++ b/tools/heir-opt.cpp
@@ -353,6 +353,15 @@ int main(int argc, char **argv) {
         mlirToSecretArithmeticPipelineBuilder(pm, options);
       });
 
+  PassPipelineRegistration<mlir::heir::PlaintextBackendOptions>(
+      "mlir-to-plaintext-backend",
+      "Convert a func using standard MLIR dialects to plaintext backend "
+      "executed in secret arithmetic",
+      [](OpPassManager &pm,
+         const mlir::heir::PlaintextBackendOptions &options) {
+        mlirToPlaintextPipelineBuilder(pm, options);
+      });
+
   PassPipelineRegistration<mlir::heir::MlirToRLWEPipelineOptions>(
       "mlir-to-bgv",
       "Convert a func using standard MLIR dialects to FHE using "


Reviews:
user: login: j2kun
body: LGTM!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-11T03:54:43Z
user: login: ZenithalHourlyRate
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T04:45:16Z
user: login: ZenithalHourlyRate
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T04:48:05Z
user: login: j2kun
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T16:52:25Z

Review comments:
diff_hunk: @@ -0,0 +1,34 @@
+// Secret arithmetic plaintext backend with original integer type semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_test.a %t -o a.out
path: tests/Examples/plaintext/dot_product/dot_product_8.mlir
user: login: j2kun
body: I see what you mean now with `cc` just being used for linking. This is fine. I think if it's easy to change this to `$CC` with an extra line in lit.cfg.py, something like `os.environ["CC"] = os.environ.get("CC", "/usr/bin/cc")`, that would be slightly preferable, but I am also happy to merge this as-is.
created_at: 2025-03-11T03:40:54Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,30 @@
+#include <stdint.h>
+#include <stdio.h>
+
+// This is the function we want to call from LLVM
+// For mod-arith, the width is 64 bits
+int64_t dot_product(
+    /* arg 0*/
+    int64_t *allocated, int64_t *aligned, int64_t offset, int64_t size,
+    int64_t stride,
+    /* arg 1*/
+    int64_t *allocated_2, int64_t *aligned_2, int64_t offset_2, int64_t size_2,
+    int64_t stride_2);
+
+int main() {
+  int64_t arg0[8] = {1, 2, 3, 4, 5, 6, 7, 8};
+  int64_t arg1[8] = {2, 3, 4, 5, 6, 7, 8, 9};
+  int64_t expected = 240;
path: tests/Examples/plaintext/dot_product/dot_product_8_mod_test.c
user: login: j2kun
body: I think it would be worth choosing a plaintext modulus that is small enough in one of the tests so that the expected value is different from the test in which no modulus is set. How about a modulus of `17`, so that the expected value is `2`?

Another thing to make this easier is to have this `main` print out `res`, and then have the test assertion happen in FileCheck using `--check-prefix=CHECK_MY_UNIQUE_PREFIX` as you do in other lit tests in this file.

[edit]: I noticed you did this for roberts cross, but having it for a smaller test might be better so regressions could be caught by smaller, faster running tests.
created_at: 2025-03-11T03:46:03Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,34 @@
+// Secret arithmetic plaintext backend with original integer type semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_test.a %t -o a.out
path: tests/Examples/plaintext/dot_product/dot_product_8.mlir
user: login: ZenithalHourlyRate
body: Seems that natively lit does not support bash-style command line in RUN

```
# .---command stderr------------
# | '$CC': command not found
# `-----------------------------
```

To fix it we have to use something like 
```
sh -c '$CC %S/libdot_product_8_test.a %t -o a.out'
```

I think that would be quite inelegant. Lets make cc here and whenever someone runs into issue we adopt this way (if they run into issue they already need to set $CC anyway so some changes are ok, afterall, they should not pay special care to some test cases that infrequently changes and could be convered in CI)
created_at: 2025-03-11T04:45:16Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,30 @@
+#include <stdint.h>
+#include <stdio.h>
+
+// This is the function we want to call from LLVM
+// For mod-arith, the width is 64 bits
+int64_t dot_product(
+    /* arg 0*/
+    int64_t *allocated, int64_t *aligned, int64_t offset, int64_t size,
+    int64_t stride,
+    /* arg 1*/
+    int64_t *allocated_2, int64_t *aligned_2, int64_t offset_2, int64_t size_2,
+    int64_t stride_2);
+
+int main() {
+  int64_t arg0[8] = {1, 2, 3, 4, 5, 6, 7, 8};
+  int64_t arg1[8] = {2, 3, 4, 5, 6, 7, 8, 9};
+  int64_t expected = 240;
path: tests/Examples/plaintext/dot_product/dot_product_8_mod_test.c
user: login: ZenithalHourlyRate
body: Added the modulus=17 test case
created_at: 2025-03-11T04:48:05Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,34 @@
+// Secret arithmetic plaintext backend with original integer type semantic
+//
+// RUN: heir-opt %s --mlir-to-plaintext-backend \
+// RUN:   | mlir-translate --mlir-to-llvmir | llc -filetype=obj > %t
+// RUN: cc %S/libdot_product_8_test.a %t -o a.out
path: tests/Examples/plaintext/dot_product/dot_product_8.mlir
user: login: j2kun
body: SGTM
created_at: 2025-03-11T16:52:25Z
author_association: COLLABORATOR

PR comments:
user: login: ZenithalHourlyRate
created_at: 2025-03-04T04:19:10Z
body: The lesson I learnt from playing with the plaintext modulus is that, even though the user input may be i16 and the output does not exceed that, after early pipeline like vectorizer/layout (which often assumes things happen in Z instead of Zt or Z_i16) the user is totally unaware of what the circuit becomes, so the user might need the help of plaintext backend to select an appropriate plaintext modulus with typical inputs. We might write a guide on selecting plaintext modulus. Note that other libraries take the same approach by either suggesting a large ptm or let the user try each one.

In such situation compiler also could not help. As one may expect, integer range analysis just gives worst-case estimate and it could be too loose. What the range analysis can do is some partial effort. The same thing should hold for the float point pipeline where range analysis is also wanted.
author_association: COLLABORATOR
user: login: j2kun
created_at: 2025-03-11T03:55:32Z
body: I'll wait until the morning to mark this as pull_ready so you have a chance to respond to the optional comments.
author_association: COLLABORATOR