Skip to content

Commit 9478d3d

Browse files
j2kuncopybara-github
authored andcommitted
Add debug dialect with validate op
This is the first step of a series of PRs that will ultimately improve the debug infrastructure and plaintext execution framework. In particular, for larger programs, the current *-add-debug-port passes insert debug functions after each op, and the plaintext execution writes the plaintext values one per line. This generates too much data, and the association of plaintext execution output to the original program points is quite brittle. This new dialect is intended to help by (a) inserting debug.validate ops at a high level, which later lower to debug function calls but does not require a debug handler after every low-level op (b) having a unique identifier for the program point at which validation occurs and (c) having a json structure for plaintext execution results (or any other metadata) that will make the plaintext dump easier to manage and relate to the original IR. PiperOrigin-RevId: 872502333
1 parent 2be4b87 commit 9478d3d

12 files changed

Lines changed: 214 additions & 0 deletions

File tree

lib/Dialect/Debug/IR/BUILD

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# Debug dialect
2+
3+
load("@heir//lib/Dialect:dialect.bzl", "add_heir_dialect_library")
4+
load("@llvm-project//mlir:tblgen.bzl", "td_library")
5+
load("@rules_cc//cc:cc_library.bzl", "cc_library")
6+
7+
package(
8+
default_applicable_licenses = ["@heir//:license"],
9+
default_visibility = ["//visibility:public"],
10+
)
11+
12+
cc_library(
13+
name = "Dialect",
14+
srcs = [
15+
"DebugDialect.cpp",
16+
"DebugOps.cpp",
17+
],
18+
hdrs = [
19+
"DebugDialect.h",
20+
"DebugOps.h",
21+
],
22+
deps = [
23+
":dialect_inc_gen",
24+
":ops_inc_gen",
25+
"@llvm-project//mlir:IR",
26+
],
27+
)
28+
29+
td_library(
30+
name = "td_files",
31+
srcs = [
32+
"DebugDialect.td",
33+
"DebugOps.td",
34+
],
35+
deps = [
36+
"@llvm-project//mlir:OpBaseTdFiles",
37+
"@llvm-project//mlir:SideEffectInterfacesTdFiles",
38+
],
39+
)
40+
41+
add_heir_dialect_library(
42+
name = "dialect_inc_gen",
43+
dialect = "Debug",
44+
kind = "dialect",
45+
td_file = "DebugDialect.td",
46+
deps = [
47+
":td_files",
48+
],
49+
)
50+
51+
add_heir_dialect_library(
52+
name = "ops_inc_gen",
53+
dialect = "Debug",
54+
kind = "op",
55+
td_file = "DebugOps.td",
56+
deps = [
57+
":td_files",
58+
],
59+
)
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#include "lib/Dialect/Debug/IR/DebugDialect.h"
2+
3+
#include "lib/Dialect/Debug/IR/DebugOps.h"
4+
5+
// Generated definitions
6+
#include "lib/Dialect/Debug/IR/DebugDialect.cpp.inc"
7+
8+
namespace mlir {
9+
namespace heir {
10+
namespace debug {
11+
12+
void DebugDialect::initialize() {
13+
addOperations<
14+
#define GET_OP_LIST
15+
#include "lib/Dialect/Debug/IR/DebugOps.cpp.inc"
16+
>();
17+
}
18+
19+
} // namespace debug
20+
} // namespace heir
21+
} // namespace mlir
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#ifndef LIB_DIALECT_DEBUG_IR_DEBUGDIALECT_H_
2+
#define LIB_DIALECT_DEBUG_IR_DEBUGDIALECT_H_
3+
4+
// IWYU pragma: begin_keep
5+
#include "mlir/include/mlir/IR/Builders.h" // from @llvm-project
6+
#include "mlir/include/mlir/IR/Dialect.h" // from @llvm-project
7+
#include "mlir/include/mlir/IR/DialectImplementation.h" // from @llvm-project
8+
// IWYU pragma: end_keep
9+
10+
// Generated headers
11+
#include "lib/Dialect/Debug/IR/DebugDialect.h.inc"
12+
13+
#endif // LIB_DIALECT_DEBUG_IR_DEBUGDIALECT_H_
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
#ifndef LIB_DIALECT_DEBUG_IR_DEBUGDIALECT_TD_
2+
#define LIB_DIALECT_DEBUG_IR_DEBUGDIALECT_TD_
3+
4+
include "mlir/IR/DialectBase.td"
5+
include "mlir/IR/OpBase.td"
6+
7+
def Debug_Dialect : Dialect {
8+
let name = "debug";
9+
10+
let description = [{
11+
The Debug dialect contains operations for debugging and validation.
12+
13+
These ops are intended to be used during development to do things like
14+
insert function calls at specific program checkpoints.
15+
}];
16+
17+
let cppNamespace = "::mlir::heir::debug";
18+
}
19+
20+
#endif // LIB_DIALECT_DEBUG_IR_DEBUGDIALECT_TD_

lib/Dialect/Debug/IR/DebugOps.cpp

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
#include "lib/Dialect/Debug/IR/DebugOps.h"
2+
3+
// IWYU pragma: begin_keep
4+
#include "mlir/include/mlir/IR/OpImplementation.h" // from @llvm-project
5+
// IWYU pragma: end_keep
6+
7+
// Generated definitions
8+
#define GET_OP_CLASSES
9+
#include "lib/Dialect/Debug/IR/DebugOps.cpp.inc"
10+
11+
namespace mlir {
12+
namespace heir {
13+
namespace debug {} // namespace debug
14+
} // namespace heir
15+
} // namespace mlir

lib/Dialect/Debug/IR/DebugOps.h

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
#ifndef LIB_DIALECT_DEBUG_IR_DEBUGOPS_H_
2+
#define LIB_DIALECT_DEBUG_IR_DEBUGOPS_H_
3+
4+
// IWYU pragma: begin_keep
5+
#include "lib/Dialect/Debug/IR/DebugDialect.h"
6+
#include "mlir/include/mlir/IR/BuiltinOps.h" // from @llvm-project
7+
#include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project
8+
#include "mlir/include/mlir/IR/Dialect.h" // from @llvm-project
9+
#include "mlir/include/mlir/IR/OpDefinition.h" // from @llvm-project
10+
// IWYU pragma: end_keep
11+
12+
#define GET_OP_CLASSES
13+
#include "lib/Dialect/Debug/IR/DebugOps.h.inc"
14+
15+
#endif // LIB_DIALECT_DEBUG_IR_DEBUGOPS_H_

lib/Dialect/Debug/IR/DebugOps.td

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#ifndef LIB_DIALECT_DEBUG_IR_DEBUGOPS_TD_
2+
#define LIB_DIALECT_DEBUG_IR_DEBUGOPS_TD_
3+
4+
include "DebugDialect.td"
5+
6+
include "mlir/IR/OpBase.td"
7+
include "mlir/Interfaces/SideEffectInterfaces.td"
8+
9+
class Debug_Op<string mnemonic, list<Trait> traits = []> :
10+
Op<Debug_Dialect, mnemonic, traits> {
11+
let cppNamespace = "::mlir::heir::debug";
12+
}
13+
14+
def Debug_ValidateOp : Debug_Op<"validate", []> {
15+
let summary = "Validates an SSA value.";
16+
let description = [{
17+
The `debug.validate` operation is a high-level placeholder for validating
18+
an SSA value. This is transformed via `*-add-debug-port` passes to a function
19+
call to an externally defined function that may then decrypt and validate
20+
the operand.
21+
22+
The mandatory `name` attribute gives a unique identifier for the validation
23+
instance, and this is used to connect intermediate values of a plaintext
24+
execution of a program to the corresponding program points of the
25+
HEIR-compiled program.
26+
27+
An optional `metadata` attribute may contain an arbitrary JSON blob, to
28+
be passed to the function call, which is intended to contain metadata like
29+
the plaintext execution result that the called function can use to compare
30+
with the decrypted ciphertext.
31+
}];
32+
33+
let arguments = (ins
34+
AnyType:$input,
35+
StrAttr:$name,
36+
OptionalAttr<StrAttr>:$metadata
37+
);
38+
39+
let assemblyFormat = "$input attr-dict `:` type($input)";
40+
}
41+
42+
#endif // LIB_DIALECT_DEBUG_IR_DEBUGOPS_TD_

tests/Dialect/Debug/IR/BUILD

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
load("//bazel:lit.bzl", "glob_lit_tests")
2+
3+
package(default_applicable_licenses = ["@heir//:license"])
4+
5+
glob_lit_tests(
6+
name = "all_tests",
7+
data = ["@heir//tests:test_utilities"],
8+
driver = "@heir//tests:run_lit.sh",
9+
test_file_exts = ["mlir"],
10+
)

tests/Dialect/Debug/IR/syntax.mlir

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// RUN: heir-opt %s | FileCheck %s
2+
3+
// CHECK: module
4+
module {
5+
// CHECK: @test_syntax
6+
func.func @test_syntax(%arg0: i32) {
7+
// CHECK: debug.validate %arg0 {name = "test_val"} : i32
8+
debug.validate %arg0 { name = "test_val" } : i32
9+
// CHECK: debug.validate %arg0 {metadata = "{\22foo\22: 7}", name = "test_val_with_metadata"} : i32
10+
debug.validate %arg0 { name = "test_val_with_metadata", metadata = "{\"foo\": 7}" } : i32
11+
return
12+
}
13+
}

tools/BUILD

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ cc_binary(
5454
"@heir//lib/Dialect/CKKS/IR:Dialect",
5555
"@heir//lib/Dialect/CKKS/Transforms",
5656
"@heir//lib/Dialect/Comb/IR:Dialect",
57+
"@heir//lib/Dialect/Debug/IR:Dialect",
5758
"@heir//lib/Dialect/Jaxite/IR:Dialect",
5859
"@heir//lib/Dialect/JaxiteWord/IR:Dialect",
5960
"@heir//lib/Dialect/LWE/Conversions/LWEToLattigo",
@@ -237,6 +238,7 @@ cc_binary(
237238
"@heir//lib/Dialect/CGGI/IR:Dialect",
238239
"@heir//lib/Dialect/CKKS/IR:Dialect",
239240
"@heir//lib/Dialect/Comb/IR:Dialect",
241+
"@heir//lib/Dialect/Debug/IR:Dialect",
240242
"@heir//lib/Dialect/Jaxite/IR:Dialect",
241243
"@heir//lib/Dialect/JaxiteWord/IR:Dialect",
242244
"@heir//lib/Dialect/LWE/IR:Dialect",

0 commit comments

Comments
 (0)