[Mlir-commits] [mlir] [mlir][affine] Add dependency on `UBDialect` for `PoisonAttr` (PR #77691)

[Felix Schneider]

https://github.com/ubfx created https://github.com/llvm/llvm-project/pull/77691

The folder for `AffineApplyOp` [1] will try creating a `PoisonAttr` under certain circumstances. However, this will result in a crash if the `UBDialect` isn't loaded:
`LLVM ERROR: can't create Attribute 'mlir::ub::PoisonAttr' because storage uniquer isn't initialized ...`

This patch adds a dependency of `AffineDialect` on `UBDialect`.

[1] https://github.com/llvm/llvm-project/blob/6d3ebd831c31d473acb18511949d04038115864a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp#L607

>From c58a040092ac716da3b4f8e47004534832fa6ce1 Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Wed, 10 Jan 2024 22:33:02 +0100
Subject: [PATCH] [mlir][affine] Add dependency on `UBDialect` for `PoisonAttr`

The folder for `AffineApplyOp` [1] will try creating a `PoisonAttr`
under certain circumstances. However, this will result in a crash
if the `UBDialect` isn't loaded:
`LLVM ERROR: can't create Attribute 'mlir::ub::PoisonAttr' because
storage uniquer isn't initialized ...`

This patch adds a dependency of `AffineDialect` on `UBDialect`.

[1] https://github.com/llvm/llvm-project/blob/6d3ebd831c31d473acb18511949d04038115864a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp#L607
---
 mlir/include/mlir/Dialect/Affine/IR/AffineOps.h  |  1 +
 mlir/include/mlir/Dialect/Affine/IR/AffineOps.td |  2 +-
 mlir/test/Dialect/Affine/constant-fold.mlir      | 11 ++++++++++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index f070d048861906..b0a14c9cae4cd2 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -16,6 +16,7 @@
 
 #include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index c638646b9c3277..225e4d3194e230 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -24,7 +24,7 @@ def Affine_Dialect : Dialect {
   let name = "affine";
   let cppNamespace = "::mlir::affine";
   let hasConstantMaterializer = 1;
-  let dependentDialects = ["arith::ArithDialect"];
+  let dependentDialects = ["arith::ArithDialect", "ub::UBDialect"];
 }
 
 // Base class for Affine dialect ops.
diff --git a/mlir/test/Dialect/Affine/constant-fold.mlir b/mlir/test/Dialect/Affine/constant-fold.mlir
index 5236b44ddfed96..ffc3946db08df5 100644
--- a/mlir/test/Dialect/Affine/constant-fold.mlir
+++ b/mlir/test/Dialect/Affine/constant-fold.mlir
@@ -64,7 +64,7 @@ func.func @affine_min(%variable: index) -> (index, index) {
 // -----
 
 func.func @affine_apply_poison_division_zero() {
-  // This is just for mlir::context to load ub dailect
+  // This is just for mlir::context to load ub dialect
   %ub = ub.poison : index
   %c16 = arith.constant 16 : index
   %0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
@@ -81,3 +81,12 @@ func.func @affine_apply_poison_division_zero() {
 // CHECK-NEXT: %[[alloc:.*]] = memref.alloc(%[[poison]], %[[poison]], %[[poison]])
 // CHECK-NEXT: %[[load:.*]] = affine.load %[[alloc]][%[[poison]], %[[poison]], %[[poison]]] : memref<?x?x?xi1>
 // CHECK-NEXT: affine.store %[[load]], %alloc[%[[poison]], %[[poison]], %[[poison]]] : memref<?x?x?xi1>
+
+// -----
+
+// Check that this doesn't crash because the ub dialect is automatically loaded
+func.func @affine_apply_poison_division_zero_2() -> index {
+  %c16 = arith.constant 16 : index
+  %0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
+  return %0 : index
+}