Pipelines

heir-opt

--heir-tosa-to-arith

Lowers a TOSA MLIR model to func, arith, and memref.

Lowers from TOSA through linalg and affine, and converts all tensors to memrefs. Fully unrolls all loops, and forwards stores to subsequent loads whenever possible. The output is suitable as an input to heir-translate --emit-verilog. Retains affine.load and affine.store ops that cannot be removed (e.g., reading from the input and writing to the output, or loading from a memref with a variable index).

The pass pipeline assumes that the input is a valid TOSA MLIR model with stripped quantized types. The iree-import-tflite tool can lower a TFLite FlatBuffer to textual MLIR with --output-format=mlir-ir. See hello_world.tosa.mlir for an example.

--yosys-optimizer

Uses Yosys to booleanize and optimize MLIR functions.

This pass pipeline requires inputs to be in standard MLIR (arith, affine, func, memref). The pass imports the model to Yosys and runs passes to booleanize the circuit and then uses ABC to perform optimizations. We use standard LUT 3 cells. THe output of this pass includes arith constants and comb.truth_table ops.

The pass requires that the environment variable HEIR_ABC_BINARY contains the location of the ABC binary and that HEIR_YOSYS_SCRIPTS_DIR contains the location of the Yosys’ techlib files that are needed to execute the path.

This pass can be disabled by defining HEIR_NO_YOSYS; this will avoid Yosys library and ABC binary compilation, and avoid registration of this pass.

heir-translate

--emit-tfhe-rust

Code generation for the tfhe-rs FHE library. The library is based on the CGGI cryptosystem, and so this pass is most useful when paired with lowerings from the cggi dialect.

The version of tfhe-rs supported is defined in the end to end tfhe_rust tests.

Example input:

!sks = !tfhe_rust.server_key
!lut = !tfhe_rust.lookup_table
!eui3 = !tfhe_rust.eui3

func.func @test_apply_lookup_table(%sks : !sks, %lut: !lut, %input : !eui3) -> !eui3 {
  %v1 = tfhe_rust.apply_lookup_table %sks, %input, %lut : (!sks, !eui3, !lut) -> !eui3
  %v2 = tfhe_rust.add %sks, %input, %v1 : (!sks, !eui3, !eui3) -> !eui3
  %c1 = arith.constant 1 : i8
  %v3 = tfhe_rust.scalar_left_shift %sks, %v2, %c1 : (!sks, !eui3, i8) -> !eui3
  %v4 = tfhe_rust.apply_lookup_table %sks, %v3, %lut : (!sks, !eui3, !lut) -> !eui3
  return %v4 : !eui3
}

Example output:

use tfhe::shortint::prelude::*;

pub fn test_apply_lookup_table(
  v9: &ServerKey,
  v10: &LookupTableOwned,
  v11: &Ciphertext,
) -> Ciphertext {
  let v4 = v9.apply_lookup_table(&v11, &v10);
  let v5 = v9.unchecked_add(&v11, &v4);
  let v6 = 1;
  let v7 = v9.scalar_left_shift(&v5, v6);
  let v8 = v9.apply_lookup_table(&v7, &v10);
  v8
}

Note, the chosen variable names are arbitrary, and the resulting program still must be integrated with a larger Rust program.

--emit-verilog

Code generation for verilog from arith and memref. Expects a single top level func.func op as the entry point, which is converted to the output verilog module.

Example input:

module {
  func.func @main(%arg0: i8) -> (i8) {
    %c0 = arith.constant 0 : i32
    %c1 = arith.constant 1 : i32
    %c2 = arith.constant 2 : i32
    %c3 = arith.constant 3 : i32
    %0 = arith.extsi %arg0 : i8 to i32
    %1 = arith.subi %0, %c1 : i32
    %2 = arith.muli %1, %c2 : i32
    %3 = arith.addi %2, %c3 : i32
    %4 = arith.cmpi sge, %2, %c0 : i32
    %5 = arith.select %4, %c1, %c2 : i32
    %6 = arith.shrsi %3, %c1 : i32
    %7 = arith.shrui %3, %c1 : i32
    %out = arith.trunci %6 : i32 to i8
    return %out : i8
  }
}

Output:

module main(
  input wire signed [7:0] arg1,
  output wire signed [7:0] _out_
);
  wire signed [31:0] v2;
  wire signed [31:0] v3;
  wire signed [31:0] v4;
  wire signed [31:0] v5;
  wire signed [31:0] v6;
  wire signed [31:0] v7;
  wire signed [31:0] v8;
  wire signed [31:0] v9;
  wire v10;
  wire signed [31:0] v11;
  wire signed [31:0] v12;
  wire signed [31:0] v13;
  wire signed [7:0] v14;

  assign v2 = 0;
  assign v3 = 1;
  assign v4 = 2;
  assign v5 = 3;
  assign v6 = {{24{arg1[7]}}, arg1};
  assign v7 = v6 - v3;
  assign v8 = v7 * v4;
  assign v9 = v8 + v5;
  assign v10 = v8 >= v2;
  assign v11 = v10 ? v3 : v4;
  assign v12 = v9 >>> v3;
  assign v13 = v9 >> v3;
  assign v14 = v12[7:0];
  assign _out_ = v14;
endmodule

--emit-metadata

Prints a json object describing the function signatures. Used for code generation after --emit-verilog.

Example input:

module {
  func.func @main(%arg0: memref<80xi8>) -> memref<1x3x2x1xi8> {
    %alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<1x3x2x1xi8>
    return %alloc_0 : memref<1x3x2x1xi8>
  }
}

Example output:

{
  "functions": [
    {
      "name": "main",
      "params": [
        {
          "index": 0,
          "type": {
            "memref": {
              "element_type": {
                "integer": {
                  "is_signed": false,
                  "width": 8
                }
              },
              "shape": [80]
            }
          }
        }
      ],
      "return_types": [{
        "memref": {
          "element_type": {
            "integer": {
              "is_signed": false,
              "width": 8
            }
          },
          "shape": [1, 3, 2, 1]
        }
      }]
    }
  ]
}