Skip to content

[mlir][affine] Add dependency on UBDialect for PoisonAttr #77691

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jan 11, 2024

Conversation

ubfx
Copy link
Member

@ubfx ubfx commented Jan 10, 2024

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]

return ub::PoisonAttr::get(getContext());

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
@llvmbot
Copy link
Member

llvmbot commented Jan 10, 2024

@llvm/pr-subscribers-mlir

Author: Felix Schneider (ubfx)

Changes

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]

return ub::PoisonAttr::get(getContext());


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

3 Files Affected:

  • (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.h (+1)
  • (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+1-1)
  • (modified) mlir/test/Dialect/Affine/constant-fold.mlir (+10-1)
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
+}

@llvmbot
Copy link
Member

llvmbot commented Jan 10, 2024

@llvm/pr-subscribers-mlir-affine

Author: Felix Schneider (ubfx)

Changes

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]

return ub::PoisonAttr::get(getContext());


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

3 Files Affected:

  • (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.h (+1)
  • (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+1-1)
  • (modified) mlir/test/Dialect/Affine/constant-fold.mlir (+10-1)
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
+}

Copy link
Member

@lipracer lipracer left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks.

@ubfx ubfx requested a review from Hardcode84 January 11, 2024 07:15
@ubfx ubfx merged commit 061b777 into llvm:main Jan 11, 2024
@ubfx ubfx deleted the fix-affine-poisonattr-storage-uniquer branch January 11, 2024 18:52
ubfx added a commit to ubfx/llvm-project that referenced this pull request Jan 11, 2024
Before applying the peeling patterns, it can happen that the `ForOp`
gets a step of zero during folding. This leads to a division-by-zero
down the line.

This patch adds an additional check for a constant-zero step and a
test.

Fix llvm#75758
Depends on llvm#77691
justinfargnoli pushed a commit to justinfargnoli/llvm-project that referenced this pull request Jan 28, 2024
…7691)

The folder for `AffineApplyOp` will try creating a `PoisonAttr`
under certain circumstances. However, this will result in a crash if the
`UBDialect` isn't loaded.

This patch adds a dependency of `AffineDialect` on `UBDialect`.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants