[Mlir-commits] [mlir] [mlir][bufferization] Check OpFilter before casting to BufferizableOpInterface (PR #85690)

[Benjamin Kramer]

https://github.com/d0k created https://github.com/llvm/llvm-project/pull/85690

This doesn't change functionality, but lets us avoid attaching all the interfaces after 513cdb82223a106f183b49a40d9acb1f7efbbe7e turned casting without loading into an error.

>From 3d0a103cdf9c8432e63772379c0fbfa7a70fa151 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <benny.kra at googlemail.com>
Date: Mon, 18 Mar 2024 20:54:02 +0100
Subject: [PATCH] [mlir][bufferization] Check OpFilter before casting to
 BufferizableOpInterface

This doesn't change functionality, but lets us avoid attaching all the
interfaces after 513cdb82223a106f183b49a40d9acb1f7efbbe7e turned casting
without loading into an error.
---
 mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp | 4 ++--
 mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp       | 4 +---
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 8f0f6d1fcc8490..55c9299c58effd 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -344,11 +344,11 @@ bool BufferizationOptions::isOpAllowed(Operation *op) const {
 
 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
+  if (!isOpAllowed(op))
+    return nullptr;
   auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
   if (!bufferizableOp)
     return nullptr;
-  if (!isOpAllowed(op))
-    return nullptr;
   return bufferizableOp;
 }
 
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 6a0ad66549965a..8dbf701620121a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -453,7 +453,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
   // canonicalize away (or canonicalize to more precise layouts).
   SmallVector<Operation *> worklist;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
-    if (hasTensorSemantics(op))
+    if (options.isOpAllowed(op) && hasTensorSemantics(op))
       worklist.push_back(op);
   });
 
@@ -472,8 +472,6 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
     auto bufferizableOp = options.dynCastBufferizableOp(nextOp);
     if (!bufferizableOp)
       continue;
-    if (!options.isOpAllowed(nextOp))
-      continue;
     // Skip ops that no longer have tensor semantics.
     if (!hasTensorSemantics(nextOp))
       continue;