[mlir][Interfaces] ValueBoundsOpInterface: Handle all destination style ops (#65736)

This commit provides a default implementation for all ops that implement
the `DestinationStyleOpInterface`. Result values of such ops are tied to
operand, and those have the same type.

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

@@ -270,24 +270,4 @@ protected:
 
 #include "mlir/Interfaces/ValueBoundsOpInterface.h.inc"
 
-namespace mlir {
-
-/// Default implementation for destination style ops: Tied OpResults and
-/// OpOperands have the same type.
-template <typename ConcreteOp>
-struct DstValueBoundsOpInterfaceExternalModel
-    : public ValueBoundsOpInterface::ExternalModel<
-          DstValueBoundsOpInterfaceExternalModel<ConcreteOp>, ConcreteOp> {
-  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
-                                       ValueBoundsConstraintSet &cstr) const {
-    auto dstOp = cast<DestinationStyleOpInterface>(op);
-    assert(value.getDefiningOp() == dstOp);
-
-    Value tiedOperand = dstOp.getTiedOpOperand(cast<OpResult>(value))->get();
-    cstr.bound(value)[dim] == cstr.getExpr(tiedOperand, dim);
-  }
-};
-
-} // namespace mlir
-
 #endif // MLIR_INTERFACES_VALUEBOUNDSOPINTERFACE_H_

mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp

@@ -47,16 +47,6 @@ struct IndexOpInterface
   }
 };
 
-/// Helper structure that iterates over all LinalgOps in `OpTys` and registers
-/// the `ValueBoundsOpInterface` with each of them.
-template <typename... Ops> struct LinalgValueBoundsOpInterfaceHelper {
-  static void registerOpInterface(MLIRContext *ctx) {
-    (Ops::template attachInterface<DstValueBoundsOpInterfaceExternalModel<Ops>>(
-         *ctx),
-     ...);
-  }
-};
-
 } // namespace
 } // namespace linalg
 } // namespace mlir
@@ -65,11 +55,7 @@ void mlir::linalg::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
     IndexOp::attachInterface<IndexOpInterface>(*ctx);
-
-    // Register all Linalg structured ops.
-    LinalgValueBoundsOpInterfaceHelper<
-#define GET_OP_LIST
-#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
-        >::registerOpInterface(ctx);
+    // Note: ValueBoundsOpInterface implementation is not required for ops that
+    // implement `DestinationStyleOpInterface` (for querying shaped OpResults).
   });
 }

mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp

@@ -120,11 +120,9 @@ void mlir::tensor::registerValueBoundsOpInterfaceExternalModels(
     tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
     tensor::ExtractSliceOp::attachInterface<tensor::ExtractSliceOpInterface>(
         *ctx);
-    tensor::InsertOp::attachInterface<
-        DstValueBoundsOpInterfaceExternalModel<tensor::InsertOp>>(*ctx);
-    tensor::InsertSliceOp::attachInterface<
-        DstValueBoundsOpInterfaceExternalModel<tensor::InsertSliceOp>>(*ctx);
     tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
     tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
+    // Note: ValueBoundsOpInterface implementation is not required for ops that
+    // implement `DestinationStyleOpInterface` (for querying shaped OpResults).
   });
 }

mlir/lib/Interfaces/CMakeLists.txt

@@ -93,6 +93,7 @@ add_mlir_library(MLIRValueBoundsOpInterface
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
 
   DEPENDS
+  MLIRDestinationStyleOpInterface
   MLIRValueBoundsOpInterfaceIncGen
 
   LINK_LIBS PUBLIC

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

@@ -10,6 +10,7 @@
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/Support/Debug.h"
 
@@ -191,13 +192,23 @@ void ValueBoundsConstraintSet::processWorklist(StopConditionFn stopCondition) {
     // the worklist.
     auto valueBoundsOp =
         dyn_cast<ValueBoundsOpInterface>(getOwnerOfValue(value));
-    if (!valueBoundsOp)
+    if (valueBoundsOp) {
+      if (dim == kIndexValue) {
+        valueBoundsOp.populateBoundsForIndexValue(value, *this);
+      } else {
+        valueBoundsOp.populateBoundsForShapedValueDim(value, dim, *this);
+      }
       continue;
-    if (dim == kIndexValue) {
-      valueBoundsOp.populateBoundsForIndexValue(value, *this);
-    } else {
-      valueBoundsOp.populateBoundsForShapedValueDim(value, dim, *this);
     }
+
+    // If the op does not implement `ValueBoundsOpInterface`, check if it
+    // implements the `DestinationStyleOpInterface`. OpResults of such ops are
+    // tied to OpOperands. Tied values have the same shape.
+    auto dstOp = value.getDefiningOp<DestinationStyleOpInterface>();
+    if (!dstOp || dim == kIndexValue)
+      continue;
+    Value tiedOperand = dstOp.getTiedOpOperand(cast<OpResult>(value))->get();
+    bound(value)[dim] == getExpr(tiedOperand, dim);
   }
 }
 

mlir/test/Dialect/Vector/value-bounds-op-interface-impl.mlir

@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \
+// RUN:     -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @vector_transfer_write(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
+//       CHECK:   return %[[dim]]
+func.func @vector_transfer_write(%t: tensor<?xf32>, %v: vector<5xf32>, %pos: index) -> index {
+  %0 = vector.transfer_write %v, %t[%pos] : vector<5xf32>, tensor<?xf32>
+  %1 = "test.reify_bound"(%0) {dim = 0} : (tensor<?xf32>) -> (index)
+  return %1 : index
+}