[Mlir-commits] [mlir] [mlir] fix crash in transform.print verification (PR #86679)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/86679

Transform op trait verification calls `getEffects`, and since trait verification runs before op verification, this call cannot assume the op to be valid. However, the operand getters now return a `TypedValue` that unconditionally casts the value to the expected type, leading to an assertion failure. Use the untyped mechanism instead.

Fixes #84701.

>From 19ab66dd1a85054cbb7eb90df793d06b2b387580 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Tue, 26 Mar 2024 15:11:59 +0000
Subject: [PATCH] [mlir] fix crash in transform.print verification

Transform op trait verification calls `getEffects`, and since trait
verification runs before op verification, this call cannot assume the op
to be valid. However, the operand getters now return a `TypedValue` that
unconditionally casts the value to the expected type, leading to an
assertion failure. Use the untyped mechanism instead.

Fixes #84701.
---
 mlir/lib/Dialect/Transform/IR/TransformOps.cpp |  6 +++++-
 mlir/test/Dialect/Transform/ops-invalid.mlir   | 11 +++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 8d2ed8f6d73714..abd557a508a103 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -2628,7 +2628,11 @@ transform::PrintOp::apply(transform::TransformRewriter &rewriter,
 
 void transform::PrintOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  onlyReadsHandle(getTarget(), effects);
+  // We don't really care about mutability here, but `getTarget` now
+  // unconditionally casts to a specific type before verification could run
+  // here.
+  if (!getTargetMutable().empty())
+    onlyReadsHandle(getTargetMutable()[0].get(), effects);
   onlyReadsPayload(effects);
 
   // There is no resource for stderr file descriptor, so just declare print
diff --git a/mlir/test/Dialect/Transform/ops-invalid.mlir b/mlir/test/Dialect/Transform/ops-invalid.mlir
index 73a5f36af92952..cc04e65420c5b7 100644
--- a/mlir/test/Dialect/Transform/ops-invalid.mlir
+++ b/mlir/test/Dialect/Transform/ops-invalid.mlir
@@ -771,3 +771,14 @@ module attributes { transform.with_named_sequence } {
     transform.yield %arg0 : !transform.any_op
   }
 }
+
+// -----
+
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @match_matmul(%entry: !transform.any_op) -> () {
+    %c3 = transform.param.constant 1 : i64 -> !transform.param<i64>
+    // expected-error @below {{op operand #0 must be TransformHandleTypeInterface instance}}
+    transform.print %c3 : !transform.param<i64>
+    transform.yield
+  }
+}