[Mlir-commits] [llvm] [mlir] [mlir][bufferization] Add a ValueBoundsOpInterface model for to_tensor (PR #215357)
Victor Perez
llvmlistbot at llvm.org
Tue Aug 11 07:46:41 PDT 2026
https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/215357
>From 68c6e631bc9adb9af79ee2e4e41a0b5366636a3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
<victor.pc.upm at gmail.com>
Date: Tue, 11 Aug 2026 07:42:28 -0700
Subject: [PATCH] [mlir][bufferization] Add ValueBoundsOpInterface models for
to_tensor/to_buffer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The bufferization dialect attaches no `ValueBoundsOpInterface` models, so a
bound query that reaches a tensor materialized from a buffer stops at the
`bufferization.to_tensor`, even though the memref side is analyzable: upstream
models `memref.alloc`, `memref.cast`, `memref.dim`, `memref.get_global` and
`memref.subview`. The mirror case has the same problem: a query that reaches a
buffer materialized from a tensor stops at the `bufferization.to_buffer`.
This matters during bufferization. An operand that was a tensor is by then a
`to_tensor` of the buffer of the producer, so every bound query through it
fails.
Attach a model to both ops. For `to_tensor`, the tensor and its source buffer
describe the same memory, so they have the same size in every dimension. For
`to_buffer`, the op carries `SameOperandsAndResultShape`, so the buffer and its
source tensor have the same size in every dimension. Both ops also accept
tensor-like and buffer-like types that are not shaped, for which no bound can
be computed.
Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
.../IR/ValueBoundsOpInterfaceImpl.h | 20 ++++++
.../Dialect/Bufferization/IR/CMakeLists.txt | 2 +
.../IR/ValueBoundsOpInterfaceImpl.cpp | 67 +++++++++++++++++++
mlir/lib/RegisterAllDialects.cpp | 2 +
.../value-bounds-op-interface-impl.mlir | 60 +++++++++++++++++
.../llvm-project-overlay/mlir/BUILD.bazel | 3 +
6 files changed, 154 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h
create mode 100644 mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp
create mode 100644 mlir/test/Dialect/Bufferization/value-bounds-op-interface-impl.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h
new file mode 100644
index 0000000000000..be8c9b3bab0ce
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h
@@ -0,0 +1,20 @@
+//===- ValueBoundsOpInterfaceImpl.h - Impl. of ValueBoundsOpInterface -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZATION_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+#define MLIR_DIALECT_BUFFERIZATION_IR_VALUEBOUNDSOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace bufferization {
+void registerValueBoundsOpInterfaceExternalModels(DialectRegistry ®istry);
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_IR_VALUEBOUNDSOPINTERFACEIMPL_H
diff --git a/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt
index 5d8f0060f2c3f..a6403a5298dab 100644
--- a/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt
@@ -7,6 +7,7 @@ add_mlir_dialect_library(MLIRBufferizationDialect
BufferViewFlowOpInterface.cpp
UnstructuredControlFlow.cpp
BufferizationTypeInterfaces.cpp
+ ValueBoundsOpInterfaceImpl.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
@@ -28,4 +29,5 @@ add_mlir_dialect_library(MLIRBufferizationDialect
MLIRSubsetOpInterface
MLIRTensorDialect
MLIRMemRefDialect
+ MLIRValueBoundsOpInterface
)
diff --git a/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp
new file mode 100644
index 0000000000000..d68e06b0593d9
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -0,0 +1,67 @@
+//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace bufferization {
+namespace {
+
+struct ToBufferOpInterface
+ : public ValueBoundsOpInterface::ExternalModel<ToBufferOpInterface,
+ ToBufferOp> {
+ void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+ ValueBoundsConstraintSet &cstr) const {
+ auto toBufferOp = cast<ToBufferOp>(op);
+ assert(value == toBufferOp.getBuffer() && "invalid value");
+
+ // The op carries `SameOperandsAndResultShape`, so the buffer and its source
+ // tensor have the same size in every dimension. The op also accepts types
+ // that are not shaped, for which no bound can be computed.
+ if (isa<ShapedType>(toBufferOp.getBuffer().getType()) &&
+ isa<ShapedType>(toBufferOp.getTensor().getType()))
+ cstr.bound(value)[dim] == cstr.getExpr(toBufferOp.getTensor(), dim);
+ }
+};
+
+struct ToTensorOpInterface
+ : public ValueBoundsOpInterface::ExternalModel<ToTensorOpInterface,
+ ToTensorOp> {
+ void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+ ValueBoundsConstraintSet &cstr) const {
+ auto toTensorOp = cast<ToTensorOp>(op);
+ assert(value == toTensorOp.getResult() && "invalid value");
+
+ // The tensor and the buffer describe the same memory, so they have the same
+ // size in every dimension. The op also accepts types that are not shaped,
+ // for which no bound can be computed.
+ if (isa<ShapedType>(toTensorOp.getResult().getType()) &&
+ isa<ShapedType>(toTensorOp.getBuffer().getType()))
+ cstr.bound(value)[dim] == cstr.getExpr(toTensorOp.getBuffer(), dim);
+ }
+};
+
+} // namespace
+} // namespace bufferization
+} // namespace mlir
+
+void mlir::bufferization::registerValueBoundsOpInterfaceExternalModels(
+ DialectRegistry ®istry) {
+ registry.addExtension(
+ +[](MLIRContext *ctx, bufferization::BufferizationDialect *dialect) {
+ bufferization::ToBufferOp::attachInterface<
+ bufferization::ToBufferOpInterface>(*ctx);
+ bufferization::ToTensorOp::attachInterface<
+ bufferization::ToTensorOpInterface>(*ctx);
+ });
+}
diff --git a/mlir/lib/RegisterAllDialects.cpp b/mlir/lib/RegisterAllDialects.cpp
index 974b5f533860a..948abc2ae4f4b 100644
--- a/mlir/lib/RegisterAllDialects.cpp
+++ b/mlir/lib/RegisterAllDialects.cpp
@@ -28,6 +28,7 @@
#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
@@ -167,6 +168,7 @@ void mlir::registerAllDialects(DialectRegistry ®istry) {
arith::registerBufferViewFlowOpInterfaceExternalModels(registry);
arith::registerShardingInterfaceExternalModels(registry);
arith::registerValueBoundsOpInterfaceExternalModels(registry);
+ bufferization::registerValueBoundsOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
registry);
builtin::registerCastOpInterfaceExternalModels(registry);
diff --git a/mlir/test/Dialect/Bufferization/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Bufferization/value-bounds-op-interface-impl.mlir
new file mode 100644
index 0000000000000..5c71b56a95de2
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/value-bounds-op-interface-impl.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics \
+// RUN: -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @to_buffer(
+// CHECK-SAME: %[[t:.*]]: tensor<?x4xf32>
+// CHECK: %[[c0:.*]] = arith.constant 0 : index
+// CHECK: %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
+// CHECK: %[[c4:.*]] = arith.constant 4 : index
+// CHECK: return %[[dim]], %[[c4]]
+func.func @to_buffer(%t: tensor<?x4xf32>) -> (index, index) {
+ %0 = bufferization.to_buffer %t : tensor<?x4xf32> to memref<?x4xf32>
+ %1 = "test.reify_bound"(%0) {dim = 0} : (memref<?x4xf32>) -> (index)
+ %2 = "test.reify_bound"(%0) {dim = 1} : (memref<?x4xf32>) -> (index)
+ return %1, %2 : index, index
+}
+
+// -----
+
+// The query goes through the op and reaches the ops that define the tensor.
+
+// CHECK-LABEL: func @to_buffer_constant(
+// CHECK: %[[c5:.*]] = arith.constant 5 : index
+// CHECK: return %[[c5]]
+func.func @to_buffer_constant() -> index {
+ %0 = tensor.empty() : tensor<5x4xf32>
+ %1 = tensor.cast %0 : tensor<5x4xf32> to tensor<?x4xf32>
+ %2 = bufferization.to_buffer %1 : tensor<?x4xf32> to memref<?x4xf32>
+ %3 = "test.reify_bound"(%2) {dim = 0, constant} : (memref<?x4xf32>) -> (index)
+ return %3 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @to_tensor(
+// CHECK-SAME: %[[m:.*]]: memref<?x4xf32>
+// CHECK: %[[c0:.*]] = arith.constant 0 : index
+// CHECK: %[[dim:.*]] = memref.dim %[[m]], %[[c0]]
+// CHECK: %[[c4:.*]] = arith.constant 4 : index
+// CHECK: return %[[dim]], %[[c4]]
+func.func @to_tensor(%m: memref<?x4xf32>) -> (index, index) {
+ %0 = bufferization.to_tensor %m : memref<?x4xf32> to tensor<?x4xf32>
+ %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?x4xf32>) -> (index)
+ %2 = "test.reify_bound"(%0) {dim = 1} : (tensor<?x4xf32>) -> (index)
+ return %1, %2 : index, index
+}
+
+// -----
+
+// The query goes through the op and reaches the ops that define the buffer.
+
+// CHECK-LABEL: func @to_tensor_constant(
+// CHECK: %[[c5:.*]] = arith.constant 5 : index
+// CHECK: return %[[c5]]
+func.func @to_tensor_constant() -> index {
+ %0 = memref.alloc() : memref<5x4xf32>
+ %1 = memref.cast %0 : memref<5x4xf32> to memref<?x4xf32>
+ %2 = bufferization.to_tensor %1 : memref<?x4xf32> to tensor<?x4xf32>
+ %3 = "test.reify_bound"(%2) {dim = 0, constant} : (tensor<?x4xf32>) -> (index)
+ return %3 : index
+}
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index fdfaa9dd2e168..0353be24d30ee 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -14159,11 +14159,13 @@ cc_library(
"lib/Dialect/Bufferization/IR/BufferizationOps.cpp",
"lib/Dialect/Bufferization/IR/BufferizationTypeInterfaces.cpp",
"lib/Dialect/Bufferization/IR/UnstructuredControlFlow.cpp",
+ "lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp",
],
hdrs = [
"include/mlir/Dialect/Bufferization/IR/Bufferization.h",
"include/mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h",
"include/mlir/Dialect/Bufferization/IR/UnstructuredControlFlow.h",
+ "include/mlir/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.h",
],
includes = ["include"],
deps = [
@@ -14191,6 +14193,7 @@ cc_library(
":SparseTensorDialect",
":SubsetOpInterface",
":TensorDialect",
+ ":ValueBoundsOpInterface",
"//llvm:Support",
],
)
More information about the Mlir-commits
mailing list