[Mlir-commits] [mlir] [mlir][vector] Add FoldTransferReadOfEmptyTensor. (PR #196599)
Erick Ochoa Lopez
llvmlistbot at llvm.org
Fri May 8 11:18:47 PDT 2026
https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/196599
When reading an empty tensor, one can fold it to the following cases:
Case 1 — fully in-bounds, no mask:
%e = tensor.empty() : tensor<128xf16>
%r = vector.transfer_read %e[%c0], %pad {in_bounds = [true]}
->
%r = ub.poison : vector<128xf16>
Case 2 — fully in-bounds, masked:
%e = tensor.empty() : tensor<128xf16>
%r = vector.transfer_read %e[%c0], %pad, %mask {in_bounds = [true]}
->
%poison = ub.poison : vector<128xf16>
%bcast = vector.broadcast %pad : f16 to vector<128xf16>
%r = arith.select %mask, %poison, %bcast
Case 3 — not fully in-bounds, no mask:
%e = tensor.empty() : tensor<100xf16>
%r = vector.transfer_read %e[%c0], %pad : tensor<100xf16>, vector<128xf16>
->
%r = vector.broadcast %pad : f16 to vector<128xf16>
Case 4 — not fully in-bounds, masked:
%e = tensor.empty() : tensor<100xf16>
%r = vector.transfer_read %e[%c0], %pad, %mask : tensor<100xf16>, vector<128xf16>
->
%r = vector.broadcast %pad : f16 to vector<128xf16>
Assisted by: Claude Opus 4.6
>From 39ad119dc3ffb11b7c123e5b74658ef8e16d8721 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Fri, 8 May 2026 14:13:56 -0400
Subject: [PATCH] [mlir][vector] Add FoldTransferReadOfEmptyTensor.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
When reading an empty tensor, one can fold it to the following cases:
Case 1 — fully in-bounds, no mask:
%e = tensor.empty() : tensor<128xf16>
%r = vector.transfer_read %e[%c0], %pad {in_bounds = [true]}
->
%r = ub.poison : vector<128xf16>
Case 2 — fully in-bounds, masked:
%e = tensor.empty() : tensor<128xf16>
%r = vector.transfer_read %e[%c0], %pad, %mask {in_bounds = [true]}
->
%poison = ub.poison : vector<128xf16>
%bcast = vector.broadcast %pad : f16 to vector<128xf16>
%r = arith.select %mask, %poison, %bcast
Case 3 — not fully in-bounds, no mask:
%e = tensor.empty() : tensor<100xf16>
%r = vector.transfer_read %e[%c0], %pad : tensor<100xf16>, vector<128xf16>
->
%r = vector.broadcast %pad : f16 to vector<128xf16>
Case 4 — not fully in-bounds, masked:
%e = tensor.empty() : tensor<100xf16>
%r = vector.transfer_read %e[%c0], %pad, %mask : tensor<100xf16>, vector<128xf16>
->
%r = vector.broadcast %pad : f16 to vector<128xf16>
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 74 +++++++++++-
mlir/test/Dialect/Vector/canonicalize.mlir | 128 +++++++++++++++++++++
2 files changed, 201 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 51be1e4431e70..f51a3b50f5101 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5529,6 +5529,77 @@ static AffineMap inverseWithUnusedDims(AffineMap map) {
}
namespace {
+
+/// Folds transfer_read(tensor.empty).
+///
+/// Since tensor.empty has unspecified contents, reading from it produces
+/// an unspecified value, which is exactly the semantics of ub.poison.
+///
+/// Case 1 — fully in-bounds, no mask:
+/// %e = tensor.empty() : tensor<128xf16>
+/// %r = vector.transfer_read %e[%c0], %pad {in_bounds = [true]}
+/// ->
+/// %r = ub.poison : vector<128xf16>
+///
+/// Case 2 — fully in-bounds, masked:
+/// %e = tensor.empty() : tensor<128xf16>
+/// %r = vector.transfer_read %e[%c0], %pad, %mask {in_bounds = [true]}
+/// ->
+/// %poison = ub.poison : vector<128xf16>
+/// %bcast = vector.broadcast %pad : f16 to vector<128xf16>
+/// %r = arith.select %mask, %poison, %bcast
+///
+/// Case 3 — not fully in-bounds (with or without mask):
+/// Out-of-bounds lanes produce pad, in-bounds lanes read unspecified
+/// contents from tensor.empty, so we may choose pad for all lanes.
+/// ->
+/// %r = vector.broadcast %pad : f16 to vector<128xf16>
+struct FoldTransferReadOfEmptyTensor : public OpRewritePattern<TransferReadOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(TransferReadOp op,
+ PatternRewriter &rewriter) const override {
+ if (!op.hasPureTensorSemantics())
+ return failure();
+
+ if (!op.getBase().getDefiningOp<tensor::EmptyOp>())
+ return failure();
+
+ if (!op.getPermutationMap().isMinorIdentity())
+ return failure();
+
+ bool fullyInBounds =
+ llvm::all_of(op.getInBoundsValues(), [](bool v) { return v; });
+ TypedValue<VectorType> mask = op.getMask();
+
+ if (mask && fullyInBounds) {
+ Value rPad = op.getPadding();
+ assert(!isa<VectorType>(rPad.getType()) &&
+ "masked transfers on vector element types are not supported; "
+ "see verifyTransferOp");
+ Value poison = ub::PoisonOp::create(rewriter, op.getLoc(), op.getType());
+ Value padVal = vector::BroadcastOp::create(rewriter, rPad.getLoc(),
+ op.getType(), rPad);
+ rewriter.replaceOpWithNewOp<arith::SelectOp>(op, mask, poison, padVal);
+ return success();
+ }
+
+ if (!mask && fullyInBounds) {
+ rewriter.replaceOp(
+ op, ub::PoisonOp::create(rewriter, op.getLoc(), op.getType()));
+ return success();
+ }
+
+ // Not fully in-bounds (with or without mask): out-of-bounds lanes
+ // produce pad, and in-bounds lanes read unspecified contents from
+ // tensor.empty, so we may choose pad for those too.
+ Value rPad = op.getPadding();
+ rewriter.replaceOp(op, vector::BroadcastOp::create(rewriter, rPad.getLoc(),
+ op.getType(), rPad));
+ return success();
+ }
+};
+
/// Store to load forwarding for transfer operations with permuation maps.
/// Even if the permutation maps are different we can still propagate the store
/// into the load if the size of the dimensions read and written match. Then we
@@ -5629,7 +5700,8 @@ struct TransferReadAfterWriteToBroadcast
void TransferReadOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<TransferReadAfterWriteToBroadcast>(context);
+ results.add<FoldTransferReadOfEmptyTensor, TransferReadAfterWriteToBroadcast>(
+ context);
}
FailureOr<std::optional<SmallVector<Value>>>
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6aa92ab79a0dd..dd51ebd4d825c 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4407,3 +4407,131 @@ func.func @no_fold_alltrue_mask_empty_body_scalar_result(
%result = vector.mask %all_true, %passthru { vector.yield %val : i32 } : vector<1xi1> -> i32
return %result : i32
}
+
+// -----
+
+// transfer_read from a memref (not tensor semantics): pattern must not fire.
+// CHECK-LABEL: func.func @negative_read_empty_not_tensor_semantics
+// CHECK: vector.transfer_read
+func.func @negative_read_empty_not_tensor_semantics(%m: memref<128xf16>) -> vector<128xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %r = vector.transfer_read %m[%c0], %cst {in_bounds = [true]}
+ : memref<128xf16>, vector<128xf16>
+ return %r : vector<128xf16>
+}
+
+// -----
+
+// transfer_read from a regular tensor (not tensor.empty): pattern must not fire.
+// CHECK-LABEL: func.func @negative_read_not_empty_tensor
+// CHECK: vector.transfer_read
+func.func @negative_read_not_empty_tensor(%t: tensor<128xf16>) -> vector<128xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %r = vector.transfer_read %t[%c0], %cst {in_bounds = [true]}
+ : tensor<128xf16>, vector<128xf16>
+ return %r : vector<128xf16>
+}
+
+// -----
+
+// transfer_read from tensor.empty with a transposing permutation map: bail.
+// CHECK-LABEL: func.func @negative_read_empty_non_identity_map
+// CHECK: tensor.empty
+// CHECK: vector.transfer_read
+func.func @negative_read_empty_non_identity_map() -> vector<64x128xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %e = tensor.empty() : tensor<128x64xf16>
+ %r = vector.transfer_read %e[%c0, %c0], %cst
+ {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>}
+ : tensor<128x64xf16>, vector<64x128xf16>
+ return %r : vector<64x128xf16>
+}
+
+// -----
+
+// Unmasked, in-bounds read from tensor.empty -> ub.poison.
+// CHECK-LABEL: func.func @fold_read_empty_unmasked_inbounds
+// CHECK-NOT: tensor.empty
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[POISON:.*]] = ub.poison : vector<128xf16>
+// CHECK: return %[[POISON]]
+func.func @fold_read_empty_unmasked_inbounds() -> vector<128xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %e = tensor.empty() : tensor<128xf16>
+ %r = vector.transfer_read %e[%c0], %cst {in_bounds = [true]}
+ : tensor<128xf16>, vector<128xf16>
+ return %r : vector<128xf16>
+}
+
+// -----
+
+// Unmasked, out-of-bounds read from tensor.empty -> broadcast(pad).
+// CHECK-LABEL: func.func @fold_read_empty_unmasked_outofbounds
+// CHECK-NOT: tensor.empty
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[PAD:.+]] = arith.constant dense<0.000000e+00> : vector<256xf16>
+// CHECK: return %[[PAD]]
+func.func @fold_read_empty_unmasked_outofbounds() -> vector<256xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %e = tensor.empty() : tensor<128xf16>
+ %r = vector.transfer_read %e[%c0], %cst
+ : tensor<128xf16>, vector<256xf16>
+ return %r : vector<256xf16>
+}
+
+// -----
+
+// Masked, in-bounds read from tensor.empty with a concrete pad value ->
+// select(mask, poison, broadcast(pad)).
+// The canonicalizer also folds select(mask, poison, x) -> x.
+// CHECK-LABEL: func.func @fold_read_empty_masked_real_pad
+// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<128xf16>
+// CHECK: return %[[CST]]
+func.func @fold_read_empty_masked_real_pad(%mask: vector<128xi1>) -> vector<128xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %e = tensor.empty() : tensor<128xf16>
+ %r = vector.transfer_read %e[%c0], %cst, %mask {in_bounds = [true]}
+ : tensor<128xf16>, vector<128xf16>
+ return %r : vector<128xf16>
+}
+
+// -----
+
+// Masked read from tensor.empty where padding is ub.poison -> just ub.poison.
+// CHECK-LABEL: func.func @fold_read_empty_masked_poison_pad
+// CHECK-NOT: tensor.empty
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: arith.select
+// CHECK: %[[POISON:.*]] = ub.poison : vector<128xf16>
+// CHECK: return %[[POISON]]
+func.func @fold_read_empty_masked_poison_pad(%mask: vector<128xi1>) -> vector<128xf16> {
+ %c0 = arith.constant 0 : index
+ %pad = ub.poison : f16
+ %e = tensor.empty() : tensor<128xf16>
+ %r = vector.transfer_read %e[%c0], %pad, %mask {in_bounds = [true]}
+ : tensor<128xf16>, vector<128xf16>
+ return %r : vector<128xf16>
+}
+
+// -----
+
+// Unmasked read from a dynamically-shaped tensor.empty -> ub.poison.
+// CHECK-LABEL: func.func @fold_read_empty_dynamic_unmasked
+// CHECK-NOT: tensor.empty
+// CHECK-NOT: vector.transfer_read
+// CHECK: %[[POISON:.*]] = ub.poison : vector<128xf16>
+// CHECK: return %[[POISON]]
+func.func @fold_read_empty_dynamic_unmasked(%sz: index) -> vector<128xf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f16
+ %e = tensor.empty(%sz) : tensor<?xf16>
+ %r = vector.transfer_read %e[%c0], %cst {in_bounds = [true]}
+ : tensor<?xf16>, vector<128xf16>
+ return %r : vector<128xf16>
+}
More information about the Mlir-commits
mailing list