[mlir][affine] Allow memref.cast in isDimOpValidSymbol

[matthias-springer]

isDimOpValidSymbol is used during the verification of affine.for ops. It is used to check if LB/UB values are valid symbols. This change adds support for memref.cast, which can be skipped over if it is a ranked -> ranked cast.

This change fixes mlir/test/Transforms/canonicalize.mlir, which used to fail when verifying the IR after each pattern application (#74270). In this test case, a pattern that folds dynamic offsets/sizes/strides to static ones is applied. This pattern inserts a trivial memref.cast that can be folded away. This folding happens after the pattern application, so the IR fails to verify after applying the offsets/sizes/strides canonicalization pattern.

Note: The verifier of affine.for violates MLIR guidelines. Only local properties of an op should be verified. The verifier should not inspect the defining ops of operands. (This would mean that constraints such as "operand is a valid affine symbol" cannot be verified.)

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

isDimOpValidSymbol is used during the verification of affine.for ops. It is used to check if LB/UB values are valid symbols. This change adds support for memref.cast, which can be skipped over if it is a ranked -> ranked cast.

This change fixes mlir/test/Transforms/canonicalize.mlir, which used to fail when verifying the IR after each pattern application (#74270). In this test case, a pattern that folds dynamic offsets/sizes/strides to static ones is applied. This pattern inserts a trivial memref.cast that can be folded away. This folding happens after the pattern application, so the IR fails to verify after applying the offsets/sizes/strides canonicalization pattern.

Note: The verifier of affine.for violates MLIR guidelines. Only local properties of an op should be verified. The verifier should not inspect the defining ops of operands. (This would mean that constraints such as "operand is a valid affine symbol" cannot be verified.)

---

Full diff: https://github.com/llvm/llvm-project/pull/74401.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+12-1)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a7fc7ddec26e6..1b85da1cf28f4 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -354,8 +354,19 @@ static bool isDimOpValidSymbol(ShapedDimOpInterface dimOp, Region *region) {
   if (!index.has_value())
     return false;
 
+  // Skip over all memref.cast ops (if any).
+  Operation *op = dimOp.getShapedValue().getDefiningOp();
+  while (auto castOp = dyn_cast<memref::CastOp>(op)) {
+    // Bail on unranked memrefs.
+    if (isa<UnrankedMemRefType>(castOp.getSource().getType()))
+      return false;
+    op = castOp.getSource().getDefiningOp();
+    if (!op)
+      return false;
+  }
+
   int64_t i = index.value();
-  return TypeSwitch<Operation *, bool>(dimOp.getShapedValue().getDefiningOp())
+  return TypeSwitch<Operation *, bool>(op)
       .Case<memref::ViewOp, memref::SubViewOp, memref::AllocOp>(
           [&](auto op) { return isMemRefSizeValidSymbol(op, i, region); })
       .Default([](Operation *) { return false; });

[matthias-springer]

Do you see a way to improve the various affine op verifiers? If we land #74270, affine transformations become even more brittle because patterns can easily produce code that no longer verifies. Would it make sense to verify things such as mlir::affine::isValidSymbol in a separate pass and not as part of the op verifier?

[ftynse]

The change by itself LGTM.

How to structure the verifiers is a complex discussion that deserves more visibility on Discourse. We have more "non-local" verifiers in various places, and it looks important to verify such things as existence of symbols referenced by calls or presence of data layout / kernel attribute on the surrounding module. So far, I have written verifiers with the assumption that (1) verifiers can look at ancestor operations (otherwise, we cannot verify, e.g., that a terminator is nested in a specific operation) and (2) verifiers can inspect the region freely as it is a part of the operation being verified (otherwise, we cannot verify traits such as single-block regions). This was based on some discussion on the order of verification. I have implemented two-stage verifications to inspect siblings: (a) an operation verifies that its ancestor has a specific condition (typically an attribute is present); (b) when that condition is true for the parent operation, it verifies the relation between child operations in its regions. An example that you'd be familiar with is the transform.with_named_sqeuence attribute the verification of which checks of absence of recursive includes.

Affine rules are a different class because they imply traversing use-def chains, so inspecting sibling operations at various levels. Not sure if they can get into a catastrophic case where the verifier assumes that another operation was verified, e.g. to APIs of this operation, but the operation is in fact invalid and hasn't been verified. If it is, that would be a good argument for changing. Abstractly, the two-staged approach described above can be used here, too. But there is a caveat: function-level values are accepted as affine, but we don't want function op interface to be aware of the affine dialect.

[joker-eph]

(2) verifiers can inspect the region freely as it is a part of the operation being verified (otherwise, we cannot verify traits such as single-block regions).

You mean: "can inspect 'nested regions' freely" right?

[matthias-springer]

(1) verifiers can look at ancestor operations (otherwise, we cannot verify, e.g., that a terminator is nested in a specific operation)

Things like these can often be verified by traits (e.g., ParentOneOf), so you could argue that this is not a violation of any rules/best practices of a hand-written verifier.

Not sure if they can get into a catastrophic case where the verifier assumes that another operation was verified, e.g. to APIs of this operation, but the operation is in fact invalid and hasn't been verified.

I haven't seen any such cases. I see this mostly as a composability issue. affine.for verifies that its operands are defined by a hard-coded list of ops from another dialect (memref dialect). There could be a canonicalization pattern in the memref dialect that replaces a memref op with another one that is not in the hard-coded list (e.g., because we forgot to add it, as shown in this example. or because its a newly added op). now the canonicalization pattern produces invalid IR because of some complex verification in another dialect.