diff --git a/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.h b/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.h
index 2f2dc676e..d03ac9704 100644
--- a/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.h
+++ b/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.h
@@ -5,6 +5,7 @@
 #include <mlir/IR/BuiltinTypes.h>
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 #include <mlir/Interfaces/ControlFlowInterfaces.h>
+#include <mlir/IR/Builders.h>
 
 
 #include "zamalang/Dialect/MidLFHE/IR/MidLFHETypes.h"
diff --git a/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.td b/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.td
index d7d5c47e4..1b866aabd 100644
--- a/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.td
+++ b/compiler/include/zamalang/Dialect/MidLFHE/IR/MidLFHEOps.td
@@ -45,7 +45,7 @@ def PBSRegion : Region<
   "pbs region needs one block with one any integer argument">;
 
 def PBSOp : MidLFHE_Op<"pbs"> {
-  let arguments = (ins CipherTextType:$x);
+  let arguments = (ins CipherTextType:$x, I32Attr:$big_n, I32Attr:$log_noise, I32Attr:$base_log, I32Attr:$level);
   let results = (outs CipherTextType:$result);
   let regions = (region PBSRegion:$region);
 }
diff --git a/compiler/tests/Dialect/MidLFHE/op_pbs.mlir b/compiler/tests/Dialect/MidLFHE/op_pbs.mlir
index 9aaea3893..25643b345 100644
--- a/compiler/tests/Dialect/MidLFHE/op_pbs.mlir
+++ b/compiler/tests/Dialect/MidLFHE/op_pbs.mlir
@@ -6,13 +6,13 @@ func @pbs_ciphertext(%arg0: !MidLFHE.ciphertext, %arg1: i32) -> !MidLFHE.ciphert
   // CHECK-NEXT: ^bb0(%[[V2:.*]]: i32):  // no predecessors
   // CHECK-NEXT:   %[[V4:.*]] = divi_unsigned %[[V2]], %arg1 : i32
   // CHECK-NEXT:   "MidLFHE.pbs_return"(%[[V4]]) : (i32) -> ()
-  // CHECK-NEXT: }) : (!MidLFHE.ciphertext) -> !MidLFHE.ciphertext
+  // CHECK-NEXT: }) {base_log = 8 : i32, big_n = 1024 : i32, level = 2 : i32, log_noise = -20 : i32} : (!MidLFHE.ciphertext) -> !MidLFHE.ciphertext
   // CHECK-NEXT: return %[[V1]] : !MidLFHE.ciphertext
   %0 = "MidLFHE.pbs"(%arg0)({
     ^bb0(%a:i32):
       %1 = std.divi_unsigned %a, %arg1 : i32
       "MidLFHE.pbs_return"(%1) : (i32) -> ()
-  }) : (!MidLFHE.ciphertext) -> !MidLFHE.ciphertext
+  }){big_n=1024: i32, log_noise=-20: i32, base_log=8 : i32, level=2 : i32} : (!MidLFHE.ciphertext) -> !MidLFHE.ciphertext
 
   return %0 : !MidLFHE.ciphertext
 }
\ No newline at end of file