[RFC] ODS-based Name Obfuscation for Private Dialects

TLDR

Obfuscate dialect, operation, attribute, type, pass names to prevent reverse-engineering of release binaries. Guarded by a CMake flag.

static constexpr StringLiteral private_dialect::MulOp::getOperationName() {
  // With MLIR_PRIVATE_NAME_OBFUSCATOR=""
  return StringLiteral("privd.mul");

  // With MLIR_PRIVATE_NAME_OBFUSCATOR="md5sum"
  return StringLiteral("_1e8d1fee65b16e634839dea9f85d3429._353942263d1bedfbe06b7bfa78226253");
}

Motivation

We are working on an MLIR-based compiler (CUDA Tile IR) that features a public-facing dialect alongside several private, internal dialects. For IP protection and security reasons, we want to avoid exposing our internal compiler design in public binaries. MLIR retains op names, attribute names, type names and pass names as strings. These may give insight about compiler internals that we are not ready to share publicly yet.

Proposal

Distinguish between public and private dialects / passes with an ODS flag a CMake option.

Public dialects are part of the public API surface. They may be parsed / printed, have CAPI support / Python bindings, have bytecode support, have user-readable verifier error messages, etc.

Private dialects have no public API surface. Should users find a way to, e.g., dump a private op (e.g., by triggering an internal error), all they see is encrypted op names.

  1. Add a new field bit isPrivate = 0; to the TableGen classes for dialects and passes CMake option MLIR_PRIVATE_DIALECTS that specifies all private dialects as a comma-separated list.
  2. Add a new CMake option MLIR_PRIVATE_NAME_OBFUSCATOR (default: empty string) that specifies an external command for obfuscating names. Test cases can use md5sum.
  3. MLIR_PRIVATE_NAME_OBFUSCATOR enables obfuscation for all private dialects and passes.
  4. If obfuscation is enabled, the TableGen code generator obfuscates dialect names, op/attribute/type mnemonics and pass names/arguments with MLIR_PRIVATE_NAME_OBFUSCATOR.
  5. If obfuscation is enabled, the TableGen code generator does not generate parse/print methods for ops in private dialects. They will print with the generic op format and with obfuscated names. (However, it is expected that such ops would never be printed in the first place.)
  6. Add a new CMake option MLIR_PRIVATE_PASSES that obfuscates names and arguments of all passes. Descriptions are stripped entirely (empty string).

The build process can now generate two binaries: an external release binary (to be shipped to customers) and an internal binary with non-obfuscated names. The internal binary can parse both external/internal dialects and provides better error messages.

Details and Consequences

  • The proposed mechanism is entirely ODS based. No changes needed to core MLIR data structures.
  • The obfuscator (MLIR_PRIVATE_NAME_OBFUSCATOR) must be deterministic and produce a unique name. The same obfuscator must be used across all translation units.
  • Name obfuscation is at compile time and one-way. I don’t have a use case where an obfuscated op name would have to be “decrypted” within release binary code.
  • Code that refers to dialects / passes with their name or ops / types / attributes with their mnemonics in the form of a string (as opposed to the C++ class name) will no longer work. E.g., you can no longer write if (op->getDialect()->getDialectNamespace() == "arith").
  • Obfuscated ops cannot be matched with transform.match.operation_name.
  • Attribute / property names of operations (not attribute mnemonics!) are excluded from this RFC because op->getAttr("name", value) is a somewhat common pattern.
  • Obfuscated dialects are not supported by the MLIR bytecode infrastructure. (Bytecode round-trips successfully within builds that share a obfuscator, but bytecode produced with one obfuscator cannot be read by a build with a different obfuscator or no obfuscator.)
  • Verifier error messages use obfuscated op / dialect names.
  • Lit tests for private dialects must be run with binaries that were built without MLIR_PRIVATE_NAME_OBFUSCATOR.

Alternatives

Operation names could potentially be stripped entirely, i.e., replaced with empty strings. That would require changes to OperationName and/or MLIRContext. For example in MLIRContext:

/// This is a mapping from operation name to the operation info describing it.
llvm::StringMap<std::unique_ptr<OperationName::Impl>> operations;

An obfuscated op name, as opposed to an empty op name, has an advantage: customers can send meaningful “encrypted” error messages back to the developer. Obfuscated names still appear in diagnostics, so a customer can paste their error message back to the developer’s internal tools, which can de-obfuscate it with a reverse lookup table built at release time.

Prototype

Proof of concept only.

This is confusing to me: whether to “strip” a dialect from its strings or not is contextual and not a property of the dialect (whether scf or llvm should be stripped or not becomes a question of the actual compiler/tool I’m assembling). Seems to me like this should be injected via a build configuration instead (a list of dialects to strip out).

We could still generate a table though?

That makes sense. I removed the isPrivate flag from the RFC and prototype in favor of a MLIR_PRIVATE_DIALECTS CMake flag.

If needed, yes. The MLIR_PRIVATE_NAME_OBFUSCATOR command could generate the table and store in a text file. It could be implemented like this.

I understand the need for this, I just don’t understand why you need to distinguish between public and private ones. Why not just obfuscate everything?

It should be trivial to reverse the MD5 of public strings (table lookup, assuming no clashes), and “impossible” to do the same on private strings. A simple (separate) tool could decode stack traces (like c++filt) and extend the decoder with a downstream extension to the table that only you have.

If you obfuscate all dialects, you couldn’t parse the input dialect anymore. You would also obfuscate error messages that are produced by the input dialect verifier. To be fair, in the case of Tile IR that may work because our compiler accepts only our custom bytecode format and no textual IR. (The verifier / error messages part would be a bit annoying…)