[mlir-c] Add structural operation equivalence (#206537)
Exposes `OperationEquivalence` through the MLIR C API so callers can compare operations structurally rather than by handle identity. Assisted by: Claude
This commit is contained in:
Maksim Levental
GitHub
parent
84afdb4771
commit
e3bde09850
3 changed files with 240 additions and 1 deletions
|
|
@ -639,9 +639,46 @@ static inline bool mlirOperationIsNull(MlirOperation op) { return !op.ptr; }
|
|||
MLIR_CAPI_EXPORTED bool mlirOperationEqual(MlirOperation op,
|
||||
MlirOperation other);
|
||||
|
||||
/// Compute a hash for the given operation.
|
||||
/// Compute a hash for the given operation. Operand and result SSA values are
|
||||
/// hashed by identity and locations are significant, so equivalent-but-distinct
|
||||
/// operations hash differently; use mlirOperationStructuralHashValue for a hash
|
||||
/// that pairs with mlirOperationIsStructurallyEquivalent.
|
||||
MLIR_CAPI_EXPORTED size_t mlirOperationHashValue(MlirOperation op);
|
||||
|
||||
/// Flags controlling structural operation equivalence and hashing. These mirror
|
||||
/// `mlir::OperationEquivalence::Flags` and may be combined with bitwise OR.
|
||||
typedef enum MlirOperationEquivalenceFlags {
|
||||
/// No flags: locations, discardable attributes, properties and
|
||||
/// commutativity are all significant.
|
||||
MLIR_OPERATION_EQUIVALENCE_NONE = 0,
|
||||
/// Ignore the locations attached to operations.
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS = 1,
|
||||
/// Ignore the discardable attributes attached to operations.
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_DISCARDABLE_ATTRS = 2,
|
||||
/// Ignore the properties attached to operations.
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_PROPERTIES = 4,
|
||||
/// Ignore commutativity, comparing operands in an order-sensitive way.
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_COMMUTATIVITY = 8,
|
||||
} MlirOperationEquivalenceFlags;
|
||||
|
||||
/// Checks whether two operations are structurally equivalent, i.e. they have
|
||||
/// the same name, attributes, operand and result types, and recursively
|
||||
/// equivalent regions. Operand equivalence is tracked structurally while
|
||||
/// recursing into regions, so operands defined inside the compared regions need
|
||||
/// not be the exact same SSA values; operands defined outside must be. `flags`
|
||||
/// is a bitwise OR of MlirOperationEquivalenceFlags values.
|
||||
MLIR_CAPI_EXPORTED bool mlirOperationIsStructurallyEquivalent(MlirOperation lhs,
|
||||
MlirOperation rhs,
|
||||
uint32_t flags);
|
||||
|
||||
/// Computes a hash for the given operation that pairs with
|
||||
/// mlirOperationIsStructurallyEquivalent: two operations that are structurally
|
||||
/// equivalent under the same `flags` hash equally. Operands are hashed by
|
||||
/// identity, results are not hashed at all, and regions do not participate in
|
||||
/// the hash. `flags` is a bitwise OR of MlirOperationEquivalenceFlags values.
|
||||
MLIR_CAPI_EXPORTED size_t mlirOperationStructuralHashValue(MlirOperation op,
|
||||
uint32_t flags);
|
||||
|
||||
/// Gets the context this operation is associated with
|
||||
MLIR_CAPI_EXPORTED MlirContext mlirOperationGetContext(MlirOperation op);
|
||||
|
||||
|
|
|
|||
|
|
@ -661,6 +661,25 @@ size_t mlirOperationHashValue(MlirOperation op) {
|
|||
return OperationEquivalence::computeHash(unwrap(op));
|
||||
}
|
||||
|
||||
/// Translates the C equivalence flags to mlir::OperationEquivalence::Flags. The
|
||||
/// enumerator values mirror each other.
|
||||
static OperationEquivalence::Flags unwrapEquivalenceFlags(uint32_t flags) {
|
||||
return static_cast<OperationEquivalence::Flags>(flags);
|
||||
}
|
||||
|
||||
bool mlirOperationIsStructurallyEquivalent(MlirOperation lhs, MlirOperation rhs,
|
||||
uint32_t flags) {
|
||||
return OperationEquivalence::isEquivalentTo(unwrap(lhs), unwrap(rhs),
|
||||
unwrapEquivalenceFlags(flags));
|
||||
}
|
||||
|
||||
size_t mlirOperationStructuralHashValue(MlirOperation op, uint32_t flags) {
|
||||
return OperationEquivalence::computeHash(
|
||||
unwrap(op), /*hashOperands=*/OperationEquivalence::directHashValue,
|
||||
/*hashResults=*/OperationEquivalence::ignoreHashValue,
|
||||
unwrapEquivalenceFlags(flags));
|
||||
}
|
||||
|
||||
MlirContext mlirOperationGetContext(MlirOperation op) {
|
||||
return wrap(unwrap(op)->getContext());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3034,6 +3034,187 @@ int testReplaceUsesWithIf(MlirContext ctx) {
|
|||
return 0;
|
||||
}
|
||||
|
||||
int testOperationEquivalence(MlirContext ctx) {
|
||||
fprintf(stderr, "@testOperationEquivalence\n");
|
||||
// CHECK-LABEL: @testOperationEquivalence
|
||||
|
||||
mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("arith"));
|
||||
|
||||
const char *moduleStr = "func.func @f(%arg0: i32) -> i32 {\n"
|
||||
" %0 = arith.constant 42 : i32\n"
|
||||
" %1 = arith.constant 42 : i32\n"
|
||||
" %2 = arith.constant 7 : i32\n"
|
||||
" %3 = arith.subi %0, %2 : i32\n"
|
||||
" %4 = arith.subi %0, %2 : i32\n"
|
||||
" %5 = arith.subi %2, %0 : i32\n"
|
||||
" %6 = arith.subi %0, %2 {dialect.discardable} : "
|
||||
"i32\n"
|
||||
" %7 = arith.addi %0, %2 : i32\n"
|
||||
" %8 = arith.addi %2, %0 : i32\n"
|
||||
" return %0 : i32\n"
|
||||
"}\n";
|
||||
MlirModule module =
|
||||
mlirModuleCreateParse(ctx, mlirStringRefCreateFromCString(moduleStr));
|
||||
|
||||
MlirBlock moduleBody = mlirModuleGetBody(module);
|
||||
MlirOperation funcOp = mlirBlockGetFirstOperation(moduleBody);
|
||||
MlirRegion funcRegion = mlirOperationGetRegion(funcOp, 0);
|
||||
MlirBlock funcBody = mlirRegionGetFirstBlock(funcRegion);
|
||||
|
||||
MlirOperation c42a = mlirBlockGetFirstOperation(funcBody);
|
||||
MlirOperation c42b = mlirOperationGetNextInBlock(c42a);
|
||||
MlirOperation c7 = mlirOperationGetNextInBlock(c42b);
|
||||
MlirOperation sub3 = mlirOperationGetNextInBlock(c7);
|
||||
MlirOperation sub4 = mlirOperationGetNextInBlock(sub3);
|
||||
MlirOperation sub5 = mlirOperationGetNextInBlock(sub4);
|
||||
MlirOperation sub6 = mlirOperationGetNextInBlock(sub5);
|
||||
MlirOperation add7 = mlirOperationGetNextInBlock(sub6);
|
||||
MlirOperation add8 = mlirOperationGetNextInBlock(add7);
|
||||
|
||||
// Two identical constants are structurally equivalent when locations are
|
||||
// ignored, even though their result SSA values and source locations differ.
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
c42a, c42b, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// Without ignoring locations they differ, since they sit on distinct lines.
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
c42a, c42b, MLIR_OPERATION_EQUIVALENCE_NONE));
|
||||
|
||||
// A constant with a different value is not equivalent.
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
c42a, c7, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// Equivalence is reflexive.
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
c42a, c42a, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// `%3 = subi %0, %2` and `%4 = subi %0, %2` use the exact same operands, so
|
||||
// they are equivalent.
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
sub3, sub4, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// `%5 = subi %2, %0` is not equivalent to `%3 = subi %0, %2`: `subi` is not
|
||||
// commutative, so operands are compared pairwise, and the very first pair
|
||||
// (`%0` = 42 vs `%2` = 7) already mismatches. See the `addi` pair below for
|
||||
// the commutative path.
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
sub3, sub5, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// A constant and a subtraction are different operations.
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
c42a, sub3, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// `%6 = subi %0, %2 {dialect.discardable}` carries a discardable attribute
|
||||
// that `%3 = subi %0, %2` does not, so the two differ unless discardable
|
||||
// attributes are ignored as well. This also exercises OR-ing flags together.
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
sub3, sub6, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
sub3, sub6,
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS |
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_DISCARDABLE_ATTRS));
|
||||
|
||||
// `%7 = addi %0, %2` and `%8 = addi %2, %0` have swapped operands, but `addi`
|
||||
// is commutative, so they are equivalent through the commutative path and
|
||||
// only differ once commutativity is explicitly ignored.
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
add7, add8, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
add7, add8,
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS |
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_COMMUTATIVITY));
|
||||
|
||||
// The `value` of arith.constant is an inherent attribute held in the
|
||||
// operation's properties, so ignoring properties makes constants with
|
||||
// different values equivalent.
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
c42a, c7,
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS |
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_PROPERTIES));
|
||||
|
||||
// The structural hash pairs with the equivalence above: operations that are
|
||||
// equivalent under a set of flags hash equally under the same flags.
|
||||
assert(mlirOperationStructuralHashValue(
|
||||
c42a, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS) ==
|
||||
mlirOperationStructuralHashValue(
|
||||
c42b, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
assert(mlirOperationStructuralHashValue(
|
||||
sub3, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS) ==
|
||||
mlirOperationStructuralHashValue(
|
||||
sub4, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
assert(mlirOperationStructuralHashValue(
|
||||
sub3, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS |
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_DISCARDABLE_ATTRS) ==
|
||||
mlirOperationStructuralHashValue(
|
||||
sub6, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS |
|
||||
MLIR_OPERATION_EQUIVALENCE_IGNORE_DISCARDABLE_ATTRS));
|
||||
// Commutative operands are folded into the hash in an order-insensitive way,
|
||||
// matching the commutative equivalence above.
|
||||
assert(mlirOperationStructuralHashValue(
|
||||
add7, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS) ==
|
||||
mlirOperationStructuralHashValue(
|
||||
add8, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// The flags are threaded through to the hash: locations and discardable
|
||||
// attributes are hashed unless the corresponding flag is set. Hash
|
||||
// inequality is not part of the contract, but these inputs do differ.
|
||||
assert(
|
||||
mlirOperationStructuralHashValue(c42a, MLIR_OPERATION_EQUIVALENCE_NONE) !=
|
||||
mlirOperationStructuralHashValue(c42b, MLIR_OPERATION_EQUIVALENCE_NONE));
|
||||
assert(mlirOperationStructuralHashValue(
|
||||
sub3, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS) !=
|
||||
mlirOperationStructuralHashValue(
|
||||
sub6, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
mlirModuleDestroy(module);
|
||||
|
||||
// Operands defined *inside* the compared regions need not be the same SSA
|
||||
// values: recursing into the regions marks the two inner constants
|
||||
// equivalent, and the `addi` operands are then matched through that mapping.
|
||||
const char *regionModuleStr = "func.func @g() {\n"
|
||||
" %r0 = scf.execute_region -> i32 {\n"
|
||||
" %c = arith.constant 1 : i32\n"
|
||||
" %d = arith.addi %c, %c : i32\n"
|
||||
" scf.yield %d : i32\n"
|
||||
" }\n"
|
||||
" %r1 = scf.execute_region -> i32 {\n"
|
||||
" %c = arith.constant 1 : i32\n"
|
||||
" %d = arith.addi %c, %c : i32\n"
|
||||
" scf.yield %d : i32\n"
|
||||
" }\n"
|
||||
" return\n"
|
||||
"}\n";
|
||||
MlirModule regionModule = mlirModuleCreateParse(
|
||||
ctx, mlirStringRefCreateFromCString(regionModuleStr));
|
||||
assert(!mlirModuleIsNull(regionModule));
|
||||
MlirBlock regionFuncBody = mlirRegionGetFirstBlock(mlirOperationGetRegion(
|
||||
mlirBlockGetFirstOperation(mlirModuleGetBody(regionModule)), 0));
|
||||
MlirOperation exec0 = mlirBlockGetFirstOperation(regionFuncBody);
|
||||
MlirOperation exec1 = mlirOperationGetNextInBlock(exec0);
|
||||
assert(mlirOperationIsStructurallyEquivalent(
|
||||
exec0, exec1, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
// The nested `addi` operations compared on their own are *not* equivalent:
|
||||
// the value mapping starts out empty, so operands defined outside the
|
||||
// compared operation must be the exact same SSA values.
|
||||
MlirBlock exec0Body =
|
||||
mlirRegionGetFirstBlock(mlirOperationGetRegion(exec0, 0));
|
||||
MlirBlock exec1Body =
|
||||
mlirRegionGetFirstBlock(mlirOperationGetRegion(exec1, 0));
|
||||
MlirOperation innerAdd0 =
|
||||
mlirOperationGetNextInBlock(mlirBlockGetFirstOperation(exec0Body));
|
||||
MlirOperation innerAdd1 =
|
||||
mlirOperationGetNextInBlock(mlirBlockGetFirstOperation(exec1Body));
|
||||
assert(!mlirOperationIsStructurallyEquivalent(
|
||||
innerAdd0, innerAdd1, MLIR_OPERATION_EQUIVALENCE_IGNORE_LOCATIONS));
|
||||
|
||||
mlirModuleDestroy(regionModule);
|
||||
|
||||
// CHECK: testOperationEquivalence: PASSED
|
||||
fprintf(stderr, "testOperationEquivalence: PASSED\n");
|
||||
return 0;
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
MlirContext ctx = mlirContextCreate();
|
||||
registerAllUpstreamDialects(ctx);
|
||||
|
|
@ -3091,6 +3272,8 @@ int main(void) {
|
|||
return 20;
|
||||
if (testReplaceUsesWithIf(ctx))
|
||||
return 21;
|
||||
if (testOperationEquivalence(ctx))
|
||||
return 22;
|
||||
|
||||
// CHECK: DESTROY MAIN CONTEXT
|
||||
// CHECK: reportResourceDelete: resource_i64_blob
|
||||
|
|
|
|||
Loading…
Reference in a new issue