[llvm] [mlir] [mlir][bufferization] Add a ValueBoundsOpInterface model for to_tensor (PR #215357)

[Andrei Golubev via llvm-commits]

================
@@ -0,0 +1,48 @@
+//===- 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 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 &registry) {
+  registry.addExtension(
+      +[](MLIRContext *ctx, bufferization::BufferizationDialect *dialect) {
+        bufferization::ToTensorOp::attachInterface<
+            bufferization::ToTensorOpInterface>(*ctx);
+      });
----------------
andrey-golubev wrote:

I'd say let's just have the two together? it's a minor difference code-wise but I feel like the two ops should really go together as they legalize different "ends" of the IR after bufferization.

https://github.com/llvm/llvm-project/pull/215357