[Mlir-commits] [mlir] [mlir] Avoid crash in mlir-query by using dyn_cast instead of cast for FileLineColLoc (PR #145049)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Denzel-Brian Budii (chios202)

<details>
<summary>Changes</summary>

If the following query is run: 

`match getUsersByPredicate(hasOpName("memref.alloc"), hasOpName("memref.dealloc"), true)`

on the IR illustrated below, it causes the program to crash.  Upon debugging, I think the underneath cause may be because traverses `linalg.matmul` facilitates `arith.addf`  and `getForwardSlice`  traverses the op's regions and blocks.

``` mlir
func.func @<!-- -->slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
  %a = memref.alloc(%arg0, %arg2) : memref<?x?xf32>
  %b = memref.alloc(%arg2, %arg1) : memref<?x?xf32>
  %c = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
  %d = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
               outs(%c : memref<?x?xf32>)
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
               outs(%d : memref<?x?xf32>)
  memref.dealloc %c : memref<?x?xf32>
  memref.dealloc %b : memref<?x?xf32>
  memref.dealloc %a : memref<?x?xf32>
  memref.dealloc %d : memref<?x?xf32>
  return
}
```

If changed to `dyn_cast` the query succeeds:
```mlir
Match #<!-- -->1:

/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:4:8: note: "root" binds here
  %a = memref.alloc(%arg0, %arg2) : memref<?x?xf32>
       ^
/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:8:3:
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
  ^
/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:10:3:
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
  ^
Match #<!-- -->2:

/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:5:8: note: "root" binds here
  %b = memref.alloc(%arg2, %arg1) : memref<?x?xf32>
       ^
/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:8:3:
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
  ^
/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:10:3:
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
  ^
Match #<!-- -->3:

/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:6:8: note: "root" binds here
  %c = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
       ^
/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:8:3:
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
  ^
Match #<!-- -->4:

/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:7:8: note: "root" binds here
  %d = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
       ^
/home/dbudii/personal/llvm-project/mlir/test/mlir-query/backward-slice-union.mlir:10:3:
  linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
  ^
4 matches.

```


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


1 Files Affected:

- (modified) mlir/lib/Query/Matcher/MatchFinder.cpp (+13-11) 


``````````diff
diff --git a/mlir/lib/Query/Matcher/MatchFinder.cpp b/mlir/lib/Query/Matcher/MatchFinder.cpp
index 1d4817e32417d..a4fe1cd549e25 100644
--- a/mlir/lib/Query/Matcher/MatchFinder.cpp
+++ b/mlir/lib/Query/Matcher/MatchFinder.cpp
@@ -38,21 +38,23 @@ MatchFinder::collectMatches(Operation *root, DynMatcher matcher) const {
 
 void MatchFinder::printMatch(llvm::raw_ostream &os, QuerySession &qs,
                              Operation *op) const {
-  auto fileLoc = cast<FileLineColLoc>(op->getLoc());
-  SMLoc smloc = qs.getSourceManager().FindLocForLineAndColumn(
-      qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
-  llvm::SMDiagnostic diag =
-      qs.getSourceManager().GetMessage(smloc, llvm::SourceMgr::DK_Note, "");
-  diag.print("", os, true, false, true);
+  if (auto fileLoc = dyn_cast<FileLineColLoc>(op->getLoc())) {
+    SMLoc smloc = qs.getSourceManager().FindLocForLineAndColumn(
+        qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
+    llvm::SMDiagnostic diag =
+        qs.getSourceManager().GetMessage(smloc, llvm::SourceMgr::DK_Note, "");
+    diag.print("", os, true, false, true);
+  }
 }
 
 void MatchFinder::printMatch(llvm::raw_ostream &os, QuerySession &qs,
                              Operation *op, const std::string &binding) const {
-  auto fileLoc = cast<FileLineColLoc>(op->getLoc());
-  auto smloc = qs.getSourceManager().FindLocForLineAndColumn(
-      qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
-  qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,
-                                     "\"" + binding + "\" binds here");
+  if (auto fileLoc = dyn_cast<FileLineColLoc>(op->getLoc())) {
+    auto smloc = qs.getSourceManager().FindLocForLineAndColumn(
+        qs.getBufferId(), fileLoc.getLine(), fileLoc.getColumn());
+    qs.getSourceManager().PrintMessage(os, smloc, llvm::SourceMgr::DK_Note,
+                                       "\"" + binding + "\" binds here");
+  }
 }
 
 std::vector<Operation *>

``````````

</details>


https://github.com/llvm/llvm-project/pull/145049