[Mlir-commits] [mlir] [mlir][tensor] fix out-of-bound index in tensor.dim (PR #85901)
Jianbang Yang
llvmlistbot at llvm.org
Mon Mar 25 04:29:41 PDT 2024
https://github.com/dynamicheart updated https://github.com/llvm/llvm-project/pull/85901
>From 784d288b9a96061be201f17f0dd473e0a79a93fa Mon Sep 17 00:00:00 2001
From: Jianbang Yang <yangjianbang112 at gmail.com>
Date: Wed, 20 Mar 2024 15:28:27 +0800
Subject: [PATCH 1/3] [mlir][tensor] fix out-of-bound index in tensor.dim
convert tensor.dim with out-of-bound index to ub.poison
Fixes: https://github.com/llvm/llvm-project/issues/70183
---
.../Bufferization/IR/BufferizationOps.cpp | 3 +++
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 9 +++++---
mlir/test/Dialect/MemRef/resolve-dim-ops.mlir | 3 +--
mlir/test/Dialect/Tensor/canonicalize.mlir | 22 +++++++++++++++++++
4 files changed, 32 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 2b226c7a1207cf..a656c812a59feb 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -333,6 +333,9 @@ struct FoldDimOfAllocTensorOp : public OpRewritePattern<tensor::DimOp> {
auto allocTensorOp = dimOp.getSource().getDefiningOp<AllocTensorOp>();
if (!allocTensorOp || !maybeConstantIndex)
return failure();
+ if (*maybeConstantIndex < 0 ||
+ *maybeConstantIndex >= allocTensorOp.getType().getRank())
+ return failure();
if (!allocTensorOp.getType().isDynamicDim(*maybeConstantIndex))
return failure();
rewriter.replaceOp(
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index dc8843aa4e1e13..2183839d8f4290 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
@@ -45,6 +46,8 @@ Operation *TensorDialect::materializeConstant(OpBuilder &builder,
if (complex::ConstantOp::isBuildableWith(value, type))
return builder.create<complex::ConstantOp>(loc, type,
llvm::cast<ArrayAttr>(value));
+ if (auto poison = dyn_cast<ub::PoisonAttr>(value))
+ return builder.create<ub::PoisonOp>(loc, type, poison);
return nullptr;
}
@@ -738,11 +741,11 @@ OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
if (!tensorType)
return {};
- // Out of bound indices produce undefined behavior but are still valid IR.
- // Don't choke on them.
+ // Fold dim to posion if the index is out of bound. Poison represents
+ // undefined behavior.
int64_t indexVal = index.getInt();
if (indexVal < 0 || indexVal >= tensorType.getRank())
- return {};
+ return ub::PoisonAttr::get(getContext());
// Fold if the shape extent along the given index is known.
if (!tensorType.isDynamicDim(index.getInt())) {
diff --git a/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir b/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir
index 18e9a9d02e1081..72ff1ca08c1faa 100644
--- a/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir
+++ b/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir
@@ -13,10 +13,9 @@ func.func @dim_out_of_bounds(%m : memref<7x8xf32>) -> index {
// -----
// CHECK-LABEL: func @dim_out_of_bounds_2(
-// CHECK-NEXT: arith.constant
+// CHECK-NEXT: ub.poison
// CHECK-NEXT: arith.constant
// CHECK-NEXT: bufferization.alloc_tensor
-// CHECK-NEXT: tensor.dim
// CHECK-NEXT: return
func.func @dim_out_of_bounds_2(%idx1 : index, %idx2 : index) -> index {
%idx = arith.constant 7 : index
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index e5374f031be553..2914c6873f5e47 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2367,3 +2367,25 @@ func.func @dim_of_reshape_undominated(%arg0: tensor<*xf32>, %arg1: tensor<?xinde
%dim = tensor.dim %reshape, %0 : tensor<*xf32>
return %dim : index
}
+
+// -----
+
+// Test case: tensor.dim(%tensor, %out_of_bound_index) is folded into ub.poison
+// CHECK-LABEL: func @dim_out_of_bounds(
+// CHECK: ub.poison
+// CHECK-NEXT: bufferization.alloc_tensor
+// CHECK-NEXT: memref.alloc
+// CHECK-NEXT: memref.cast
+// CHECK-NEXT: affine.vector_load
+// CHECK-NEXT: return
+func.func @dim_out_of_bounds() -> vector<7xi32> {
+ %c1 = arith.constant 1 : index
+ %idx28 = index.constant 28
+ %c29 = arith.constant 29 : index
+ %3 = bufferization.alloc_tensor(%c29) : tensor<?xi16>
+ %dim = tensor.dim %3, %idx28 : tensor<?xi16>
+ %alloc_21 = memref.alloc(%c29) : memref<?x26x2xi32>
+ %16 = affine.vector_load %alloc_21[%c1, %c1, %dim] : memref<?x26x2xi32>, vector<7xi32>
+ return %16 : vector<7xi32>
+}
+
>From b07521a7a2d62ed669131221b757d06d0502cd16 Mon Sep 17 00:00:00 2001
From: Jianbang Yang <yangjianbang112 at gmail.com>
Date: Wed, 20 Mar 2024 17:13:48 +0800
Subject: [PATCH 2/3] fix typo and cmake dependency, refine mlir test
---
mlir/lib/Dialect/Tensor/IR/CMakeLists.txt | 1 +
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 2 +-
mlir/test/Dialect/Tensor/canonicalize.mlir | 4 ++--
3 files changed, 4 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
index 549b9f10388bd4..d5f805a8a6316a 100644
--- a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
@@ -34,6 +34,7 @@ add_mlir_dialect_library(MLIRTensorDialect
MLIRShapedOpInterfaces
MLIRSideEffectInterfaces
MLIRSupport
+ MLIRUBDialect
MLIRValueBoundsOpInterface
MLIRViewLikeInterface
)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 2183839d8f4290..97d17f61c815ad 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -741,7 +741,7 @@ OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
if (!tensorType)
return {};
- // Fold dim to posion if the index is out of bound. Poison represents
+ // Fold dim to poison if the index is out of bound. Poison represents
// undefined behavior.
int64_t indexVal = index.getInt();
if (indexVal < 0 || indexVal >= tensorType.getRank())
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 2914c6873f5e47..3caefa89944f21 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2372,11 +2372,11 @@ func.func @dim_of_reshape_undominated(%arg0: tensor<*xf32>, %arg1: tensor<?xinde
// Test case: tensor.dim(%tensor, %out_of_bound_index) is folded into ub.poison
// CHECK-LABEL: func @dim_out_of_bounds(
-// CHECK: ub.poison
+// CHECK: %[[poison:.*]] = ub.poison
// CHECK-NEXT: bufferization.alloc_tensor
// CHECK-NEXT: memref.alloc
// CHECK-NEXT: memref.cast
-// CHECK-NEXT: affine.vector_load
+// CHECK-NEXT: affine.vector_load %{{.*}}[{{.*}}, {{.*}}, symbol(%[[poison]])]
// CHECK-NEXT: return
func.func @dim_out_of_bounds() -> vector<7xi32> {
%c1 = arith.constant 1 : index
>From e37adc8c86db00b1fc8f75954480062011657a91 Mon Sep 17 00:00:00 2001
From: Jianbang Yang <yangjianbang112 at gmail.com>
Date: Mon, 25 Mar 2024 19:29:06 +0800
Subject: [PATCH 3/3] revert ub.posion folding
---
mlir/lib/Dialect/Tensor/IR/CMakeLists.txt | 1 -
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 9 +++------
mlir/test/Dialect/MemRef/resolve-dim-ops.mlir | 3 ++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 7 ++++---
4 files changed, 9 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
index d5f805a8a6316a..549b9f10388bd4 100644
--- a/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/IR/CMakeLists.txt
@@ -34,7 +34,6 @@ add_mlir_dialect_library(MLIRTensorDialect
MLIRShapedOpInterfaces
MLIRSideEffectInterfaces
MLIRSupport
- MLIRUBDialect
MLIRValueBoundsOpInterface
MLIRViewLikeInterface
)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 97d17f61c815ad..dc8843aa4e1e13 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -11,7 +11,6 @@
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
@@ -46,8 +45,6 @@ Operation *TensorDialect::materializeConstant(OpBuilder &builder,
if (complex::ConstantOp::isBuildableWith(value, type))
return builder.create<complex::ConstantOp>(loc, type,
llvm::cast<ArrayAttr>(value));
- if (auto poison = dyn_cast<ub::PoisonAttr>(value))
- return builder.create<ub::PoisonOp>(loc, type, poison);
return nullptr;
}
@@ -741,11 +738,11 @@ OpFoldResult DimOp::fold(FoldAdaptor adaptor) {
if (!tensorType)
return {};
- // Fold dim to poison if the index is out of bound. Poison represents
- // undefined behavior.
+ // Out of bound indices produce undefined behavior but are still valid IR.
+ // Don't choke on them.
int64_t indexVal = index.getInt();
if (indexVal < 0 || indexVal >= tensorType.getRank())
- return ub::PoisonAttr::get(getContext());
+ return {};
// Fold if the shape extent along the given index is known.
if (!tensorType.isDynamicDim(index.getInt())) {
diff --git a/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir b/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir
index 72ff1ca08c1faa..18e9a9d02e1081 100644
--- a/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir
+++ b/mlir/test/Dialect/MemRef/resolve-dim-ops.mlir
@@ -13,9 +13,10 @@ func.func @dim_out_of_bounds(%m : memref<7x8xf32>) -> index {
// -----
// CHECK-LABEL: func @dim_out_of_bounds_2(
-// CHECK-NEXT: ub.poison
+// CHECK-NEXT: arith.constant
// CHECK-NEXT: arith.constant
// CHECK-NEXT: bufferization.alloc_tensor
+// CHECK-NEXT: tensor.dim
// CHECK-NEXT: return
func.func @dim_out_of_bounds_2(%idx1 : index, %idx2 : index) -> index {
%idx = arith.constant 7 : index
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 3caefa89944f21..9ab54fe9c133db 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2370,13 +2370,14 @@ func.func @dim_of_reshape_undominated(%arg0: tensor<*xf32>, %arg1: tensor<?xinde
// -----
-// Test case: tensor.dim(%tensor, %out_of_bound_index) is folded into ub.poison
+// Test case: This test fails to fold because the index of tensor.dim is out_of_bounds
// CHECK-LABEL: func @dim_out_of_bounds(
-// CHECK: %[[poison:.*]] = ub.poison
+// CHECK: %[[IDX:.*]] = index.constant 28
// CHECK-NEXT: bufferization.alloc_tensor
+// CHECK-NEXT: %[[DIM:.*]] = tensor.dim %{{.*}}, %[[IDX]]
// CHECK-NEXT: memref.alloc
// CHECK-NEXT: memref.cast
-// CHECK-NEXT: affine.vector_load %{{.*}}[{{.*}}, {{.*}}, symbol(%[[poison]])]
+// CHECK-NEXT: affine.vector_load %{{.*}}[{{.*}}, {{.*}}, symbol(%[[DIM]])]
// CHECK-NEXT: return
func.func @dim_out_of_bounds() -> vector<7xi32> {
%c1 = arith.constant 1 : index
More information about the Mlir-commits
mailing list