fix copybara reversibility issue

FUTURE_COPYBARA_INTEGRATE_REVIEW=#270 from j2kun:generalize-yosys-optimizer 9e56415
PiperOrigin-RevId: 586718990

include/Dialect/Secret/IR/SecretOps.td

@@ -25,7 +25,7 @@ def Secret_ConcealOp : Secret_Op<"conceal", [Pure]> {
 
     Examples:
 
-    ```
+    ```mlir
     %Y = secret.conceal %value : i32 -> !secret.secret<i32>
     ```
   }];
@@ -55,7 +55,7 @@ def Secret_RevealOp : Secret_Op<"reveal", [Pure]> {
 
     Examples:
 
-    ```
+    ```mlir
     %Y = secret.reveal %secret_value : !secret.secret<i32> -> i32
     ```
   }];
@@ -104,7 +104,7 @@ def Secret_GenericOp : Secret_Op<"generic", [
 
     Add two secret integers together
 
-    ```
+    ```mlir
     %Z = secret.generic ins(%X, %Y : !secret.secret<i32>, !secret.secret<i32>) {
       ^bb0(%x: i32, %y: i32) :
         %z = arith.addi %x, %y: i32
@@ -115,7 +115,7 @@ def Secret_GenericOp : Secret_Op<"generic", [
     Add a secret value with a plaintext value. I.e., not all arguments to the
     op need be secret.
 
-    ```
+    ```mlir
     %Z = secret.generic ins(%X, %Y : i32, !secret.secret<i32>) {
       ^bb0(%x: i32, %y: i32) :
         %z = arith.addi %x, %y: i32
@@ -154,5 +154,35 @@ def Secret_GenericOp : Secret_Op<"generic", [
   let hasVerifier = 1;
 }
 
+def CastOp : Secret_Op<"cast", [Pure]> {
+  let summary = "A placeholder cast from one secret type to another";
+  let description = [{
+    A `cast` operation represents a type cast from one secret type to another,
+    that is used to enable the intermixing of various equivalent secret types
+    before a lower-level FHE scheme has been chosen.
+
+    For example, `secret.cast` can be used to convert a `secret<i8>` to a
+    `secret<tensor<8xi1>>` as a compatibility layer between boolean and
+    non-boolean parts of a program. The pass that later lowers the IR to
+    specific FHE schemes would need to replace these casts with appropriate
+    scheme-specific operations, and it is left to those later passes to
+    determine which casts are considered valid.
+
+    Example:
+
+    ```mlir
+    %result = secret.cast %0 : !secret.secret<i8> to !secret.secret<tensor<8xi1>>
+    %result2 = secret.cast %0 : !secret.secret<i8> to !secret.secret<tensor<2xi4>>
+    ```
+  }];
+
+  let arguments = (ins Secret:$input);
+  let results = (outs Secret:$output);
+  let assemblyFormat = [{
+    $input attr-dict `:` qualified(type($input)) `to` qualified(type($output))
+  }];
+  let hasFolder = 1;
+}
+
 
 #endif  // HEIR_INCLUDE_DIALECT_SECRET_IR_SECRETOPS_TD_

include/Target/Verilog/BUILD

@@ -12,6 +12,7 @@ cc_library(
     name = "verilog_emitter",
     hdrs = ["VerilogEmitter.h"],
     deps = [
+        "@heir//lib/Dialect/Secret/IR:Dialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:ArithDialect",

include/Target/Verilog/VerilogEmitter.h

@@ -1,16 +1,29 @@
 #ifndef HEIR_INCLUDE_TARGET_VERILOG_VERILOGEMITTER_H_
 #define HEIR_INCLUDE_TARGET_VERILOG_VERILOGEMITTER_H_
 
-#include "llvm/include/llvm/ADT/DenseMap.h"    // from @llvm-project
-#include "llvm/include/llvm/ADT/TypeSwitch.h"  // from @llvm-project
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <string_view>
+
+#include "include/Dialect/Secret/IR/SecretOps.h"
+#include "llvm/include/llvm/ADT/DenseMap.h"         // from @llvm-project
+#include "llvm/include/llvm/ADT/StringRef.h"        // from @llvm-project
+#include "llvm/include/llvm/ADT/ilist.h"            // from @llvm-project
+#include "llvm/include/llvm/Support/raw_ostream.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
 #include "mlir/include/mlir/Dialect/Math/IR/Math.h"      // from @llvm-project
 #include "mlir/include/mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinOps.h"             // from @llvm-project
 #include "mlir/include/mlir/IR/Operation.h"              // from @llvm-project
+#include "mlir/include/mlir/IR/Region.h"                 // from @llvm-project
+#include "mlir/include/mlir/IR/TypeRange.h"              // from @llvm-project
+#include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
 #include "mlir/include/mlir/Support/IndentedOstream.h"   // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
 #include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project
 
 namespace mlir {
@@ -22,11 +35,26 @@ void registerToVerilogTranslation();
 mlir::LogicalResult translateToVerilog(mlir::Operation *op,
                                        llvm::raw_ostream &os);
 
+/// Translates the given operation to Verilog with a fixed input name for the
+/// resulting verilog module. Raises an error if the input IR contains secret
+/// ops.
+mlir::LogicalResult translateToVerilog(
+    mlir::Operation *op, llvm::raw_ostream &os,
+    std::optional<llvm::StringRef> moduleName);
+
+/// Translates the given operation to Verilog with a fixed input name for the
+/// resulting verilog module. If allowSecretOps is false, raises an error if
+/// the input IR contains secret ops.
+mlir::LogicalResult translateToVerilog(
+    mlir::Operation *op, llvm::raw_ostream &os,
+    std::optional<llvm::StringRef> moduleName, bool allowSecretOps);
+
 class VerilogEmitter {
  public:
   VerilogEmitter(raw_ostream &os);
 
-  LogicalResult translate(mlir::Operation &operation);
+  LogicalResult translate(mlir::Operation &operation,
+                          std::optional<llvm::StringRef> moduleName);
 
  private:
   /// Output stream to emit to.
@@ -38,8 +66,26 @@ class VerilogEmitter {
   // Globally unique identifiers for values
   int64_t value_count_;
 
+  // A helper to generalize the work of emitting a func.func and a
+  // secret.generic
+  LogicalResult printFunctionLikeOp(Operation *op,
+                                    llvm::StringRef verilogModuleName,
+                                    ArrayRef<BlockArgument> arguments,
+                                    TypeRange resultTypes,
+                                    Region::BlockListType::iterator blocksBegin,
+                                    Region::BlockListType::iterator blocksEnd);
+
+  // A helper to generalize the work of emitting a func.return and a
+  // secret.yield
+  LogicalResult printReturnLikeOp(Value returnValue);
+
   // Functions for printing individual ops
-  LogicalResult printOperation(mlir::ModuleOp op);
+  LogicalResult printOperation(mlir::ModuleOp op,
+                               std::optional<llvm::StringRef> moduleName);
+  LogicalResult printOperation(mlir::func::FuncOp op,
+                               std::optional<llvm::StringRef> moduleName);
+  LogicalResult printOperation(mlir::heir::secret::GenericOp op,
+                               std::optional<llvm::StringRef> moduleName);
   LogicalResult printOperation(mlir::UnrealizedConversionCastOp op);
   LogicalResult printOperation(mlir::arith::AddIOp op);
   LogicalResult printOperation(mlir::arith::AndIOp op);
@@ -59,8 +105,8 @@ class VerilogEmitter {
   LogicalResult printOperation(mlir::affine::AffineLoadOp op);
   LogicalResult printOperation(mlir::affine::AffineStoreOp op);
   LogicalResult printOperation(mlir::func::CallOp op);
-  LogicalResult printOperation(mlir::func::FuncOp op);
   LogicalResult printOperation(mlir::func::ReturnOp op);
+  LogicalResult printOperation(mlir::heir::secret::YieldOp op);
   LogicalResult printOperation(mlir::math::CountLeadingZerosOp op);
   LogicalResult printOperation(mlir::memref::LoadOp op);
 
@@ -69,7 +115,7 @@ class VerilogEmitter {
                               mlir::Value rhs, std::string_view op);
 
   // Emit a Verilog type of the form `wire [width-1:0]`
-  LogicalResult emitType(Location loc, Type type);
+  LogicalResult emitType(Type type);
 
   // Emit a Verilog array shape specifier of the form `[width]`
   LogicalResult emitArrayShapeSuffix(Type type);

include/Transforms/YosysOptimizer/YosysOptimizer.h

@@ -21,8 +21,8 @@ struct YosysOptimizerPipelineOptions
 // registerYosysOptimizerPipeline registers a Yosys pipeline pass using
 // runfiles, the location of Yosys techlib files, and abcPath, the location of
 // the abc binary.
-void registerYosysOptimizerPipeline(std::string yosysFilesPath,
-                                    std::string abcPath);
+void registerYosysOptimizerPipeline(const std::string &yosysFilesPath,
+                                    const std::string &abcPath);
 
 }  // namespace heir
 }  // namespace mlir

include/Transforms/YosysOptimizer/YosysOptimizer.td

@@ -18,6 +18,7 @@ def YosysOptimizer : Pass<"yosys-optimizer"> {
   let dependentDialects = [
     "mlir::arith::ArithDialect",
     "mlir::heir::comb::CombDialect",
+    "mlir::heir::secret::SecretDialect",
     "mlir::tensor::TensorDialect"
   ];
 }

lib/Dialect/Secret/IR/SecretOps.cpp

@@ -234,6 +234,28 @@ void GenericOp::build(OpBuilder &builder, OperationState &result,
   bodyBuilder(builder, result.location, bodyBlock.getArguments());
 }
 
+OpFoldResult CastOp::fold(CastOp::FoldAdaptor adaptor) {
+  Value input = getInput();
+  Value output = getOutput();
+
+  // self cast is a no-op
+  if (input.getType() == output.getType()) {
+    return input;
+  }
+
+  // Fold a cast-and-cast-back to a no-op.
+  //
+  //  %1 = secret.cast %0 : !secret.secret<T> to !secret.secret<U>
+  //  %2 = secret.cast %1 : !secret.secret<U> to !secret.secret<T>
+  //
+  // folds to use %0 directly in place of %2.
+  auto inputOp = input.getDefiningOp<CastOp>();
+  if (!inputOp || output.getType() != inputOp.getInput().getType())
+    return OpFoldResult();
+
+  return inputOp.getInput();
+}
+
 }  // namespace secret
 }  // namespace heir
 }  // namespace mlir

lib/Target/Verilog/BUILD

@@ -12,6 +12,7 @@ cc_library(
     deps = [
         "@heir//include/Target/Verilog:verilog_emitter",
         "@heir//lib/Conversion/MemrefToArith:Utils",
+        "@heir//lib/Dialect/Secret/IR:Dialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineAnalysis",
         "@llvm-project//mlir:AffineDialect",