[mlir][Interfaces] Add ReifyShapeDimTypeInterface

[springerm]

Summary

This interface allows callers to reify dimension sizes of a shaped value without depending on any dialect. This interface and its helper functions can replace switch-case statements that generate tensor.dim/memref.dim ops (depending on the type of the shaped value) in various parts of the code base.

This revision also removes createOrFoldDimOp and similiar functions from the Linalg dialect and the Vector dialect.

[aartbik]

LG for sparse

[springerm]

@mehdi_amini Just took another look at your original comment: Can we replace this with an interface method on ShapedTypeInterface and avoid hard-depending on specific dialect implementations?

This revision is adding a new interface instead of adding the functionality to ShapedTypeInterface. The problem with ShapedTypeInterface is that it would then have to be implemented as an external model in the tensor/memref dialect. This means:

• All code that uses ShapedType will have to depend on TensorDialect/MemRefDialect.
• Calls to getRank, getShaped etc. through the interface may get a bit more expensive.

I'm not sure which is better, the implementation in this revision or making ShapedTypeInterface an external model. It would be nice if we could make this work without adding a new interface.

[springerm]

We are looking into moving shaped types into dialects. Putting this revision on hold.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 bool isElementwise(LinalgOp op);
 
 /// Check if iterator type has "parallel" semantics.
 bool isParallelIterator(utils::IteratorType iteratorType);
 
 /// Check if iterator type  has "reduction" semantics.
 bool isReductionIterator(utils::IteratorType iteratorType);
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
-OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
-                               int64_t dim);
-
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
 
 /// Computes an upper bound for the result `value` of an index computation.
 /// Translates AffineMinOps and AffineApplyOps along the use-def chains of the
 /// index computation to affine constraints and projects out intermediate
 /// values. The method sets `boundMap` to an affine map that given

mlir/include/mlir/Dialect/MemRef/IR/ReifyShapeDimTypeInterfaceImpl.h

+//===- ReifyShapeDimTypeInterfaceImpl.h - Interface Implementation --------===//
+//
+// 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_MEMREF_IR_REIFYSHAPEDIMTYPEINTERFACEIMPL_H
+#define MLIR_DIALECT_MEMREF_IR_REIFYSHAPEDIMTYPEINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace memref {
+void registerReifyShapeDimTypeInterfaceExternalModels(
+    DialectRegistry &registry);
+} // namespace memref
+} // namespace mlir
+
+#endif // MLIR_DIALECT_MEMREF_IR_REIFYSHAPEDIMTYPEINTERFACEIMPL_H

mlir/include/mlir/Dialect/Tensor/IR/ReifyShapeDimTypeInterfaceImpl.h

+//===- ReifyShapeDimTypeInterfaceImpl.h - Interface Implementation --------===//
+//
+// 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_TENSOR_IR_REIFYSHAPEDIMTYPEINTERFACEIMPL_H
+#define MLIR_DIALECT_TENSOR_IR_REIFYSHAPEDIMTYPEINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace tensor {
+void registerReifyShapeDimTypeInterfaceExternalModels(
+    DialectRegistry &registry);
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_IR_REIFYSHAPEDIMTYPEINTERFACEIMPL_H

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

 class Location;
 class OpBuilder;
 class Operation;
 class ShapedType;
 class Value;
 class VectorType;
 class VectorTransferOpInterface;
 
-namespace vector {
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
-} // namespace vector
-
 /// Constructs a permutation map of invariant memref indices to vector
 /// dimension.
 ///
 /// If no index is found to be invariant, 0 is added to the permutation_map and
 /// corresponds to a vector broadcast along that dimension.
 ///
 /// The implementation uses the knowledge of the mapping of loops to
 /// vector dimension. `loopToVectorDim` carries this information as a map with:

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/IR/ReifyShapeDimTypeInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
 #include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/ReifyShapeDimTypeInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
inline void registerAllDialects(DialectRegistry &registry) {
 
   // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
+  memref::registerReifyShapeDimTypeInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
+  tensor::registerReifyShapeDimTypeInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);
 }
 
 /// Append all the MLIR dialects to the registry contained in the given context.
 inline void registerAllDialects(MLIRContext &context) {
   DialectRegistry registry;
   registerAllDialects(registry);
   context.appendDialectRegistry(registry);
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(CastInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(DestinationStyleOpInterface)
 add_mlir_interface(InferIntRangeInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
 add_mlir_interface(ParallelCombiningOpInterface)
 add_mlir_interface(RuntimeVerifiableOpInterface)
 add_mlir_interface(ShapedOpInterfaces)
+add_mlir_interface(ShapedTypeInterfaces)
 add_mlir_interface(SideEffectInterfaces)
 add_mlir_interface(TilingInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 
 set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)
 mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls)
 mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs)

mlir/include/mlir/Interfaces/ShapedTypeInterfaces.h

+//===- ShapedTypeInterfaces.h - Interfaces for Shaped Types -----*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a set of interfaces for shaped types.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_SHAPEDTYPEINTERFACES_H_
+#define MLIR_INTERFACES_SHAPEDTYPEINTERFACES_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+/// Include the generated interface declarations.
+#include "mlir/Interfaces/ShapedTypeInterfaces.h.inc"
+
+namespace mlir {
+
+/// Reify the specified dimension of the given shaped value. This function
+/// returns an IntegerAttr if `dim` is a constant and the respective dimension
+/// is static. Otherwise, it returns a Value.
+OpFoldResult reifyShapeDim(OpBuilder &builder, Location loc, Value shapedValue,
+                           OpFoldResult dim);
+
+/// Reify the specified dimension of the given shaped value. This function
+/// returns an IntegerAttr if the dimension size is static. Otherwise, it
+/// returns a Value.
+OpFoldResult reifyShapeDim(OpBuilder &builder, Location loc, Value shapedValue,
+                           int64_t dim);
+
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_SHAPEDTYPEINTERFACES_H_

mlir/include/mlir/Interfaces/ShapedTypeInterfaces.td

+//===-- ShapedTypeInterfaces.td - Interf. for Shaped Types -*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a set of interfaces for shaped types.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_SHAPEDTYPEINTERFACES
+#define MLIR_INTERFACES_SHAPEDTYPEINTERFACES
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// ReifyShapeDimTypeInterface
+//===----------------------------------------------------------------------===//
+
+def ReifyShapeDimTypeInterface : TypeInterface<"ReifyShapeDimTypeInterface"> {
+  let cppNamespace = "::mlir";
+  let description = [{
+    This is an interface for shaped types to reify dimension sizes. A reified
+    dynamic dimension is typically the result of a `DimOp` in the respective
+    dialect.
+  }];
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Reify the specified dimension of the given shaped value. If the
+        dimension is static, an `IntegerAttr` should be returned. Otherwise, a
+        `Value` should be returned.
+      }],
+      /*retTy=*/"::mlir::OpFoldResult",
+      /*methodName=*/"reifyShapeDim",
+      /*args=*/(ins "::mlir::OpBuilder &":$builder,
+                    "::mlir::Location":$loc,
+                    "::mlir::Value":$shapedValue,
+                    "::mlir::OpFoldResult":$dim)
+    >,
+  ];
+}
+
+#endif // MLIR_INTERFACES_SHAPEDTYPEINTERFACES

mlir/lib/Conversion/VectorToSCF/CMakeLists.txt

 add_mlir_conversion_library(MLIRVectorToSCF
   VectorToSCF.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/VectorToSCF
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRLLVMDialect
   MLIRMemRefDialect
+  MLIRShapedTypeInterfaces
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorTransforms
   )

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTVECTORTOSCF
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
static Value generateInBoundsCheck(
   bool hasRetVal = !resultTypes.empty();
   Value cond; // Condition to be built...
 
   // Condition check 1: Access in-bounds?
   bool isBroadcast = !dim; // No in-bounds check for broadcasts.
   Location loc = xferOp.getLoc();
   ImplicitLocOpBuilder lb(xferOp.getLoc(), b);
   if (!xferOp.isDimInBounds(0) && !isBroadcast) {
-    Value memrefDim =
-        vector::createOrFoldDimOp(b, loc, xferOp.getSource(), *dim);
+    Value memrefDim = getValueOrCreateConstantIndexOp(
+        b, loc, reifyShapeDim(b, loc, xferOp.getSource(), *dim));
     AffineExpr d0, d1;
     bindDims(xferOp.getContext(), d0, d1);
     Value base = xferOp.getIndices()[*dim];
     Value memrefIdx = makeComposedAffineApply(b, loc, d0 + d1, {base, iv});
     cond = lb.create<arith::CmpIOp>(arith::CmpIPredicate::sgt, memrefDim,
                                     memrefIdx);
   }
 

mlir/lib/Dialect/Linalg/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgDialect
   MLIRArithDialect
   MLIRArithUtils
   MLIRBufferizationDialect
   MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser
+  MLIRShapedTypeInterfaces
   MLIRSideEffectInterfaces
   MLIRSparseTensorDialect
   MLIRSCFDialect
   MLIRMathDialect
   MLIRMemRefDialect
   MLIRTensorDialect
   MLIRTilingInterface
   MLIRViewLikeInterface
   )

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

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 /// Include the definitions of the copy operation interface.
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.cpp.inc"
 
LogicalResult mlir::linalg::detail::verifyFillInterface(Operation *op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // StructuredOpInterface implementation
 //===----------------------------------------------------------------------===//
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-static Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
-                               int64_t dim) {
-  if (source.getType().isa<UnrankedMemRefType, MemRefType>())
-    return b.createOrFold<memref::DimOp>(loc, source, dim);
-  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
-    return b.createOrFold<tensor::DimOp>(loc, source, dim);
-  llvm_unreachable("Expected MemRefType or TensorType");
-}
-static OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
-                                      int64_t dim) {
-  auto shapedType = source.getType().cast<ShapedType>();
-  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
-    return createOrFoldDimOp(b, loc, source, dim);
-  return b.getIndexAttr(shapedType.getDimSize(dim));
-}
-
 SmallVector<OpFoldResult> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
                                                                 Location loc) {
   SmallVector<OpFoldResult> res;
   for (OpOperand &opOperand : getOperation()->getOpOperands()) {
     for (int64_t i = 0, e = getRank(&opOperand); i < e; ++i)
-      res.push_back(createFoldedDimOp(b, loc, opOperand.get(), i));
+      res.push_back(reifyShapeDim(b, loc, opOperand.get(), i));
   }
   return res;
 }
 
 SmallVector<int64_t, 4> LinalgOp::createFlatListOfOperandStaticDims() {
   SmallVector<int64_t, 4> res;
   assert(!hasDynamicShape() && "expected operands to have static shapes");
   for (OpOperand &opOperand : getOperation()->getOpOperands())
for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) {
       auto shapedType = opOperand->get().getType().cast<ShapedType>();
       if (!shapedType.isDynamicDim(dim)) {
         // Static dim: Return IntegerAttr.
         shapes.push_back(b.getIndexAttr(shapedType.getDimSize(dim)));
       } else {
         // Dynamic dim: Return Value.
         OpFoldResult ofr =
             checkDimExpr.visit(shapeExprs[pos])
-                ? createOrFoldDimOp(b, loc, opOperand->get(), dim)
+                ? getValueOrCreateConstantIndexOp(
+                      b, loc, reifyShapeDim(b, loc, opOperand->get(), dim))
                 : allResultDimValues[pos];
         shapes.push_back(getValueOrCreateConstantIndexOp(b, loc, ofr));
       }
       pos++;
     }
     reifiedReturnShapes.emplace_back(std::move(shapes));
   }
   return success();

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRComplexDialect
   MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRFuncDialect
   MLIRFuncToLLVM
   MLIRFuncTransforms
   MLIRInferTypeOpInterface
   MLIRIR
-  MLIRMemRefDialect
-  MLIRMemRefTransforms
   MLIRLinalgDialect
   MLIRLinalgUtils
+  MLIRMemRefDialect
+  MLIRMemRefTransforms
+  MLIRPass
   MLIRSCFDialect
   MLIRSCFTransforms
   MLIRSCFUtils
-  MLIRPass
+  MLIRShapedTypeInterfaces
   MLIRSparseTensorDialect
   MLIRTensorDialect
   MLIRTensorTilingInterfaceImpl
   MLIRTensorTransforms
   MLIRTransforms
   MLIRTransformUtils
   MLIRVectorDialect
   MLIRVectorTransforms
   MLIRVectorUtils
   MLIRX86VectorTransforms
   MLIRVectorToSCF
 )

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
static LinalgOp fuse(OpBuilder &b, LinalgOp producer,
                      const DenseMap<unsigned, Range> &fusedLoopsAndRanges) {
   SmallVector<OpFoldResult> ivs, tileSizes, sizeBounds;
   SmallVector<Range> loopRanges;
   Location loc = producer.getLoc();
 
   for (unsigned i = 0, e = producer.getNumLoops(); i < e; ++i) {
     auto shapeDim = getShapeDefiningLoopRange(producer, i);
     OpFoldResult dim =
-        createFoldedDimOp(b, loc, shapeDim.shape, shapeDim.dimension);
+        reifyShapeDim(b, loc, shapeDim.shape, shapeDim.dimension);
     sizeBounds.push_back(dim);
     auto it = fusedLoopsAndRanges.find(i);
     if (it != fusedLoopsAndRanges.end()) {
       ivs.push_back(it->second.offset);
       tileSizes.push_back(it->second.size);
       loopRanges.push_back(it->second);
       LLVM_DEBUG(llvm::dbgs() << "tiled loop#" << i << " with LoopRange "
                               << loopRanges.back() << "\n");

mlir/lib/Dialect/Linalg/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgUtils
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineAnalysis
   MLIRAffineUtils
   MLIRArithDialect
   MLIRFuncDialect
   MLIRIR
   MLIRLinalgDialect
   MLIRSCFDialect
+  MLIRShapedTypeInterfaces
   MLIRPass
   MLIRTensorUtils
   MLIRTransformUtils
   )

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "linalg-utils"
 
 using namespace mlir;

 bool isParallelIterator(utils::IteratorType iteratorType) {
   return iteratorType == utils::IteratorType::parallel;
 }
 
 bool isReductionIterator(utils::IteratorType iteratorType) {
   return iteratorType == utils::IteratorType::reduction;
 }
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
-  if (source.getType().isa<UnrankedMemRefType, MemRefType>())
-    return b.createOrFold<memref::DimOp>(loc, source, dim);
-  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
-    return b.createOrFold<tensor::DimOp>(loc, source, dim);
-  llvm_unreachable("Expected MemRefType or TensorType");
-}
-
-OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
-                               int64_t dim) {
-  auto shapedType = source.getType().cast<ShapedType>();
-  if (!shapedType.hasRank() || shapedType.isDynamicDim(dim))
-    return createOrFoldDimOp(b, loc, source, dim);
-  return b.getIndexAttr(shapedType.getDimSize(dim));
-}
-
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
   SmallVector<Value, 4> dynOperands;
   auto shapedType = val.getType().cast<ShapedType>();
   for (const auto &dim : llvm::enumerate(shapedType.getShape())) {
     if (dim.value() == ShapedType::kDynamic)
-      dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
+      dynOperands.push_back(getValueOrCreateConstantIndexOp(
+          b, loc, reifyShapeDim(b, loc, val, dim.index())));
   }
   return dynOperands;
 }
 
 void getUpperBoundForIndex(Value value, AffineMap &boundMap,
                            SmallVectorImpl<Value> &boundOperands,
                            bool constantRequired) {
   // Initialize `boundMap` and `boundOperands` to the identity returning
computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
   SliceParameters sliceParams;
   sliceParams.offsets.reserve(rank);
   sliceParams.sizes.reserve(rank);
   sliceParams.strides.reserve(rank);
   for (unsigned r = 0; r < rank; ++r) {
     LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: for dim#" << r);
     if (!isTiled(map.getSubMap({r}), tileSizes)) {
       sliceParams.offsets.push_back(builder.getIndexAttr(0));
-      OpFoldResult dim = createFoldedDimOp(builder, loc, valueToTile, r);
+      OpFoldResult dim = reifyShapeDim(builder, loc, valueToTile, r);
       sliceParams.sizes.push_back(dim);
       sliceParams.strides.push_back(builder.getIndexAttr(1));
       LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
       continue;
     }
     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
 
     // Tiling creates a new slice at the proper index, the slice step is 1

mlir/lib/Dialect/MemRef/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefDialect
   MemRefDialect.cpp
   MemRefOps.cpp
+  ReifyShapeDimTypeInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
 
   DEPENDS
   MLIRMemRefOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRArithUtils
   MLIRControlFlowInterfaces
   MLIRDialect
   MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRShapedOpInterfaces
+  MLIRShapedTypeInterfaces
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
 )

mlir/lib/Dialect/MemRef/IR/ReifyShapeDimTypeInterfaceImpl.cpp

+//===- ReifyShapeDimTypeInterfaceImpl.cpp - Interface Implementation ------===//
+//
+// 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/MemRef/IR/ReifyShapeDimTypeInterfaceImpl.h"
+
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
+
+using namespace mlir;
+using namespace mlir::memref;
+
+namespace mlir {
+namespace memref {
+namespace {
+
+template <typename T>
+struct MemRefTypeInterface
+    : public ReifyShapeDimTypeInterface::ExternalModel<MemRefTypeInterface<T>,
+                                                       T> {
+  OpFoldResult reifyShapeDim(Type type, OpBuilder &builder, Location loc,
+                             Value shapedValue, OpFoldResult dim) const {
+    BaseMemRefType memrefType = type.cast<BaseMemRefType>();
+
+    // Static dimension size: Return OpFoldResult.
+    auto constDim = getConstantIntValue(dim);
+    if (constDim.has_value() && memrefType.hasRank() &&
+        !memrefType.isDynamicDim(*constDim))
+      return builder.getIndexAttr(memrefType.getDimSize(*constDim));
+
+    // Dynamic dimension size: Create tensor.dim op.
+    return builder.createOrFold<memref::DimOp>(
+        loc, shapedValue, getValueOrCreateConstantIndexOp(builder, loc, dim));
+  }
+};
+
+} // namespace
+} // namespace memref
+} // namespace mlir
+
+void mlir::memref::registerReifyShapeDimTypeInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+    MemRefType::attachInterface<MemRefTypeInterface<MemRefType>>(*ctx);
+    UnrankedMemRefType::attachInterface<
+        MemRefTypeInterface<UnrankedMemRefType>>(*ctx);
+  });
+}

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRSparseTensorTransforms
   MLIRLLVMDialect
   MLIRLinalgDialect
   MLIRLinalgTransforms
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
   MLIRSCFTransforms
   MLIRSCFUtils
+  MLIRShapedTypeInterfaces
   MLIRSparseTensorDialect
   MLIRSparseTensorEnums
   MLIRSparseTensorUtils
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
 )

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 /// If the tensor is a sparse constant, generates and returns the pair of
 /// the constants for the coordinates and the values.
 static std::optional<std::pair<Value, Value>>
void mlir::sparse_tensor::genDenseTensorOrSparseConstantIterLoop(
   SmallVector<Value> st;
   const Value zero = constantIndex(builder, loc, 0);
   const Value one = constantIndex(builder, loc, 1);
   const auto splitSrc = genSplitSparseConstant(builder, loc, src);
   if (splitSrc.has_value()) {
     const Value srcCoordinates = splitSrc->first;
     const Value srcValues = splitSrc->second;
     lo.push_back(zero);
-    hi.push_back(linalg::createOrFoldDimOp(builder, loc, srcValues, 0));
+    hi.push_back(getValueOrCreateConstantIndexOp(
+        builder, loc, reifyShapeDim(builder, loc, srcValues, 0)));
     st.push_back(one);
     scf::buildLoopNest(builder, loc, lo, hi, st, {},
                        [&](OpBuilder &builder, Location loc, ValueRange ivs,
                            ValueRange /*args*/) -> scf::ValueVector {
                          Value val = genCoordsAndValueForSparse(
                              builder, loc, srcCoordinates, srcValues, cvs,
                              ivs[0], rank);
                          bodyBuilder(builder, loc, val, cvs);
                          return {};
                        });
   } else {
     for (unsigned i = 0; i < rank; i++) {
       lo.push_back(zero);
-      hi.push_back(linalg::createOrFoldDimOp(builder, loc, src, i));
+      hi.push_back(getValueOrCreateConstantIndexOp(
+          builder, loc, reifyShapeDim(builder, loc, src, i)));
       st.push_back(one);
     }
     scf::buildLoopNest(builder, loc, lo, hi, st, {},
                        [&](OpBuilder &builder, Location loc, ValueRange ivs,
                            ValueRange /*args*/) -> scf::ValueVector {
                          Value val = genCoordsAndValueForDense(builder, loc,
                                                                src, cvs, ivs);
                          bodyBuilder(builder, loc, val, cvs);
                          return {};
                        });
   }
 }
 
 void mlir::sparse_tensor::sizesFromSrc(OpBuilder &builder,
                                        SmallVectorImpl<Value> &sizes,
                                        Location loc, Value src) {
   const Dimension dimRank = getSparseTensorType(src).getDimRank();
   for (Dimension d = 0; d < dimRank; d++)
-    sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, d));
+    sizes.push_back(getValueOrCreateConstantIndexOp(
+        builder, loc, reifyShapeDim(builder, loc, src, d)));
 }
 
 Operation *mlir::sparse_tensor::getTop(Operation *op) {
   for (; isa<scf::ForOp>(op->getParentOp()) ||
          isa<scf::WhileOp>(op->getParentOp()) ||
          isa<scf::ParallelOp>(op->getParentOp()) ||
          isa<scf::IfOp>(op->getParentOp());
        op = op->getParentOp())

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp

 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // File local helper functions.
 //===----------------------------------------------------------------------===//
 
for (Level l = 0; l < lvlRank; l++) {
       } else {
         // Dense level, nothing to fetch.
         assert(isDenseDLT(dlt));
       }
 
       // FIXME: `toOrigDim` is deprecated
       // Since we do not have HigherOrdering now, we can always rely on the 1:1
       // mapping from level to dimension to retrieve the level size.
-      Value lvlSz = mlir::linalg::createOrFoldDimOp(builder, loc, tensor,
-                                                    toOrigDim(enc, l));
+      Value lvlSz = getValueOrCreateConstantIndexOp(
+          builder, loc, reifyShapeDim(builder, loc, tensor, toOrigDim(enc, l)));
       // Find upper bound in current dimension.
       highs[t][l] = lvlSizes[t][l] = lvlSz;
       if (isSparseSlices[t]) {
         sliceOffsets[t][l] = genSliceOffset(builder, loc, tensors[t], l);
         sliceStrides[t][l] = genSliceStride(builder, loc, tensors[t], l);
       }
     }
 

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {

 /// Returns a memref that fits the requested length (reallocates if requested
 /// length is larger, or creates a subview if it is smaller).
 static Value reallocOrSubView(OpBuilder &builder, Location loc, int64_t len,
                               Value buffer) {
   MemRefType memTp = getMemRefType(buffer);
   auto retTp = MemRefType::get(ArrayRef{len}, memTp.getElementType());
 
   Value targetLen = constantIndex(builder, loc, len);
-  Value bufferLen = linalg::createOrFoldDimOp(builder, loc, buffer, 0);
+  Value bufferLen = getValueOrCreateConstantIndexOp(
+      builder, loc, reifyShapeDim(builder, loc, buffer, 0));
   // Reallocates if target length is greater than the actual buffer len.
   Value reallocP = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ugt,
                                                  targetLen, bufferLen);
   scf::IfOp ifOp = builder.create<scf::IfOp>(loc, retTp, reallocP, true);
   // If targetLen > bufferLen, reallocate to get enough sparse to return.
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   Value reallocBuf = builder.create<memref::ReallocOp>(loc, retTp, buffer);
   builder.create<scf::YieldOp>(loc, reallocBuf);
foreachFieldAndTypeInSparseTensor(
           if (fKind == SparseTensorFieldKind::StorageSpec) {
             fields.push_back(srcDesc.getSpecifier());
           } else {
             // Allocates new memrefs
             Value srcMem = srcDesc.getMemRefField(fIdx);
             // TODO: We can instead use the actual memSize in specifier, that
             // would require a subViewOp to avoid overflow when copying
             // values.
-            Value sz = linalg::createOrFoldDimOp(rewriter, loc, srcMem, 0);
+            Value sz = getValueOrCreateConstantIndexOp(
+                rewriter, loc, reifyShapeDim(rewriter, loc, srcMem, 0));
             auto dstMem = rewriter.create<memref::AllocOp>(
                 loc, fTp.cast<MemRefType>(), sz);
             if (fTp != srcMem.getType()) {
               // Converts elements type.
               scf::buildLoopNest(
                   rewriter, loc, constantIndex(rewriter, loc, 0), sz,
                   constantIndex(rewriter, loc, 1),
                   [srcMem, &dstMem](OpBuilder &builder, Location loc,
foreachFieldAndTypeInSparseTensor(
           if (fType != field.getType())
             field = rewriter.create<memref::CastOp>(loc, fType, field);
           fields.push_back(field);
           // Returns true to continue the iteration.
           return true;
         });
 
     MutSparseTensorDescriptor desc(rtp, fields);
-    auto noe = linalg::createOrFoldDimOp(rewriter, loc, op.getValues(), 0);
+    auto noe = getValueOrCreateConstantIndexOp(
+        rewriter, loc, reifyShapeDim(rewriter, loc, op.getValues(), 0));
     // FIXME: should use `SparseTensorType::getLvlRank` in lieu of
     // `RankedTensorType::getRank`, because the latter introduces dim/lvl
     // ambiguity.
     for (Level lvl = 0, lvlRank = rtp.getRank(); lvl < lvlRank; lvl++) {
       const auto sh = rtp.getShape()[lvl];
       assert(!ShapedType::isDynamic(sh));
       desc.setLvlSize(rewriter, loc, lvl, constantIndex(rewriter, loc, sh));
       if (lvl == 0)

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
static Value createOrFoldLvlCall(OpBuilder &builder, Location loc,
   // must dynamically query it.  (In principle we could also dynamically
   // compute it, but since we already did so to construct the `tensor`
   // in the first place, we might as well query rather than recompute.)
   return genLvlSizeCall(builder, loc, tensor, lvl);
 }
 
 /// Looks up a dimension-size by returning a constant from the shape
 /// (for static sizes), or by calling `genDimSizeCall` (for dynamic sizes
-/// of sparse tensors) or `linalg::createOrFoldDimOp` (for dynamic sizes
-/// of dense tensors).
+/// of sparse tensors) or `reifyShapeDim` (for dynamic sizes of dense tensors).
 static Value createOrFoldDimCall(OpBuilder &builder, Location loc,
                                  SparseTensorType stt, Value tensor,
                                  Dimension dim) {
   if (const auto sz = stt.getStaticDimSize(dim))
     return constantIndex(builder, loc, *sz);
   if (stt.hasEncoding())
     return genDimSizeCall(builder, loc, tensor, dim);
-  return linalg::createOrFoldDimOp(builder, loc, tensor, dim);
+  return getValueOrCreateConstantIndexOp(
+      builder, loc, reifyShapeDim(builder, loc, tensor, dim));
 }
 
 /// Populates the array with the dimension-sizes of the given tensor.
 static void fillDimSizes(OpBuilder &builder, Location loc, SparseTensorType stt,
                          Value tensor, SmallVectorImpl<Value> &out) {
   const Dimension dimRank = stt.getDimRank();
   out.clear();
   out.reserve(dimRank);
static void genDenseTensorIterationLoop(
 
   SmallVector<Value> lo;
   SmallVector<Value> hi;
   SmallVector<Value> st;
 
   // Fill out loop iteration information.
   for (Dimension d = 0; d < dimRank; d++) {
     lo.push_back(zero);
-    hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, t, d));
+    hi.push_back(getValueOrCreateConstantIndexOp(
+        rewriter, loc, reifyShapeDim(rewriter, loc, t, d)));
     st.push_back(one);
   }
 
   scf::buildLoopNest(rewriter, loc, lo, hi, st, {},
                      [&](OpBuilder &builder, Location loc, ValueRange ivs,
                          ValueRange args) -> scf::ValueVector {
                        // Invoke callback to build the body of the loop.
                        bodyBuilder(builder, loc, ivs);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::linalg;
 using namespace mlir::sparse_tensor;
 
 //===---------------------------------------------------------------------===//
static void concatSizesFromInputs(OpBuilder &builder,
 
   // Sum up on the `dim` if the dimension is dynamic.
   if (dstShape[dim] != ShapedType::kDynamic) {
     // Faithfully take the static size.
     sizes[dim] = constantIndex(builder, loc, dstShape[dim]);
   } else {
     // Else, compute the shape dynamically.
     for (const auto &src : srcs.drop_front()) {
-      Value srcSz = linalg::createOrFoldDimOp(builder, loc, src, dim);
+      Value srcSz = getValueOrCreateConstantIndexOp(
+          builder, loc, reifyShapeDim(builder, loc, src, dim));
       // Sum up all the sizes.
       sizes[dim] = builder.create<arith::AddIOp>(loc, sizes[dim], srcSz);
     }
   }
 }
 
 //===---------------------------------------------------------------------===//
 // The actual sparse tensor rewriting rules.

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   TensorDialect.cpp
   TensorInferTypeOpInterfaceImpl.cpp
   TensorOps.cpp
   TensorTilingInterfaceImpl.cpp
 )
 
 add_mlir_dialect_library(MLIRTensorDialect
+  ReifyShapeDimTypeInterfaceImpl.cpp
   TensorDialect.cpp
   TensorOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/mlir/Dialect/Tensor
 
   DEPENDS
   MLIRTensorOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRArithDialect
   MLIRArithUtils
   MLIRCastInterfaces
   MLIRComplexDialect
   MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRParallelCombiningOpInterface
   MLIRShapedOpInterfaces
+  MLIRShapedTypeInterfaces
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRViewLikeInterface
   )
 
 add_mlir_dialect_library(MLIRTensorInferTypeOpInterfaceImpl
   TensorInferTypeOpInterfaceImpl.cpp
 

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

+//===- ReifyShapeDimTypeInterfaceImpl.cpp - Interface Implementation ------===//
+//
+// 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/Tensor/IR/ReifyShapeDimTypeInterfaceImpl.h"
+
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace mlir {
+namespace tensor {
+namespace {
+
+template <typename T>
+struct TensorTypeInterface
+    : public ReifyShapeDimTypeInterface::ExternalModel<TensorTypeInterface<T>,
+                                                       T> {
+  OpFoldResult reifyShapeDim(Type type, OpBuilder &builder, Location loc,
+                             Value shapedValue, OpFoldResult dim) const {
+    TensorType tensorType = type.cast<TensorType>();
+
+    // Static dimension size: Return OpFoldResult.
+    auto constDim = getConstantIntValue(dim);
+    if (constDim.has_value() && tensorType.hasRank() &&
+        !tensorType.isDynamicDim(*constDim))
+      return builder.getIndexAttr(tensorType.getDimSize(*constDim));
+
+    // Dynamic dimension size: Create tensor.dim op.
+    return builder.createOrFold<tensor::DimOp>(
+        loc, shapedValue, getValueOrCreateConstantIndexOp(builder, loc, dim));
+  }
+};
+
+} // namespace
+} // namespace tensor
+} // namespace mlir
+
+void mlir::tensor::registerReifyShapeDimTypeInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+    RankedTensorType::attachInterface<TensorTypeInterface<RankedTensorType>>(
+        *ctx);
+    UnrankedTensorType::attachInterface<
+        TensorTypeInterface<UnrankedTensorType>>(*ctx);
+  });
+}

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRTensorTransforms
   MLIRArithDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect
   MLIRPass
   MLIRSCFDialect
+  MLIRShapedTypeInterfaces
   MLIRTensorDialect
   MLIRTilingInterface
   MLIRTransforms
 )

mlir/lib/Dialect/Tensor/Transforms/ExtractSliceFromReshapeUtils.cpp

 //===- ExtractSliceFromReshapeUtils.cpp - Slice reshape rewrites ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements rewrites that replace slices of reshape results with
 // aggregated slices of the reshape source.
 //
 //===----------------------------------------------------------------------===//
+
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/TransformUtils.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
-/// Get the dimension size of a value of RankedTensor type at the
-static OpFoldResult getShapeDimSize(OpBuilder &b, Location loc,
-                                    Value rankedTensor, int64_t dimIdx) {
-  RankedTensorType tensorType = rankedTensor.getType().cast<RankedTensorType>();
-  if (!tensorType.isDynamicDim(dimIdx)) {
-    return b.getIndexAttr(tensorType.getDimSize(dimIdx));
-  }
-  Value idxValue = b.create<arith::ConstantIndexOp>(loc, dimIdx);
-  return b.createOrFold<tensor::DimOp>(loc, rankedTensor, idxValue);
-}
-
 /// Get all the dimension sizes of a value of RankedTensor type.
 static SmallVector<OpFoldResult> getShapeDimSizes(OpBuilder &b, Location loc,
                                                   Value rankedTensor) {
   SmallVector<OpFoldResult> dimSizes;
   RankedTensorType tensorType = rankedTensor.getType().cast<RankedTensorType>();
   for (unsigned i = 0; i < tensorType.getRank(); i++)
-    dimSizes.push_back(getShapeDimSize(b, loc, rankedTensor, i));
+    dimSizes.push_back(reifyShapeDim(b, loc, rankedTensor, i));
   return dimSizes;
 }
 
 /// A tuple that represents (dimension number, dimension value).
 using DimAndIndex = std::tuple<unsigned, Value>;
 
 /// Transform `dimAndIndex` from the output index space of a (non-rank-reducing)
 /// slice described by `sliceParams` into the input index space.

mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRVectorTransforms
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRDialectUtils
   MLIRGPUOps
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect
   MLIRSCFDialect
+  MLIRShapedTypeInterfaces
   MLIRSideEffectInterfaces
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorInterfaces
   MLIRVectorUtils
   )

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

 #include <optional>
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
-
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
     // Fold or create the check that `index + vector_size` <= `memref_size`.
     Location loc = xferOp.getLoc();
     int64_t vectorSize = xferOp.getVectorType().getDimSize(resultIdx);
     auto d0 = getAffineDimExpr(0, xferOp.getContext());
     auto vs = getAffineConstantExpr(vectorSize, xferOp.getContext());
     Value sum =
         makeComposedAffineApply(b, loc, d0 + vs, xferOp.indices()[indicesIdx]);
     Value cond = createFoldedSLE(
-        b, sum, vector::createOrFoldDimOp(b, loc, xferOp.source(), indicesIdx));
+        b, sum,
+        getValueOrCreateConstantIndexOp(
+            b, loc, reifyShapeDim(b, loc, xferOp.source(), indicesIdx)));
     if (!cond)
       return;
     // Conjunction over all dims for which we are in-bounds.
     if (inBoundsCond)
       inBoundsCond = b.create<arith::AndIOp>(loc, inBoundsCond, cond);
     else
       inBoundsCond = cond;
   });
assert(memrefRank == alloc.getType().cast<MemRefType>().getRank() &&
          "Expected memref rank to match the alloc rank");
   ValueRange leadingIndices =
       xferOp.indices().take_front(xferOp.getLeadingShapedRank());
   SmallVector<OpFoldResult, 4> sizes;
   sizes.append(leadingIndices.begin(), leadingIndices.end());
   auto isaWrite = isa<vector::TransferWriteOp>(xferOp);
   xferOp.zipResultAndIndexing([&](int64_t resultIdx, int64_t indicesIdx) {
     using MapList = ArrayRef<ArrayRef<AffineExpr>>;
-    Value dimMemRef = vector::createOrFoldDimOp(b, xferOp.getLoc(),
-                                                xferOp.source(), indicesIdx);
+    Value dimMemRef = getValueOrCreateConstantIndexOp(
+        b, xferOp.getLoc(),
+        reifyShapeDim(b, xferOp.getLoc(), xferOp.source(), indicesIdx));
     Value dimAlloc = b.create<memref::DimOp>(loc, alloc, resultIdx);
     Value index = xferOp.indices()[indicesIdx];
     AffineExpr i, j, k;
     bindDims(xferOp.getContext(), i, j, k);
     SmallVector<AffineMap, 4> maps =
         AffineMap::inferFromExprList(MapList{{i - j, k}});
     // affine_min(%dimMemRef - %index, %dimAlloc)
     Value affineMin = b.create<AffineMinOp>(

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/ShapedTypeInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
LogicalResult matchAndRewrite(ConcreteOp xferOp,
 
     // Create the in-bounds mask with all elements between [0 .. dim - offset)
     // set and [dim - offset .. vector_length) unset.
     //
     // TODO: when the leaf transfer rank is k > 1, we need the last `k`
     //       dimensions here.
     unsigned lastIndex = llvm::size(xferOp.getIndices()) - 1;
     Value off = xferOp.getIndices()[lastIndex];
-    Value dim =
-        vector::createOrFoldDimOp(rewriter, loc, xferOp.getSource(), lastIndex);
+    Value dim = getValueOrCreateConstantIndexOp(
+        rewriter, loc,
+        reifyShapeDim(rewriter, loc, xferOp.getSource(), lastIndex));
     Value b = rewriter.create<arith::SubIOp>(loc, dim.getType(), dim, off);
     Value mask = rewriter.create<vector::CreateMaskOp>(
         loc,
         VectorType::get(vtp.getShape(), rewriter.getI1Type(),
                         vtp.getNumScalableDims()),
         b);
     if (xferOp.getMask()) {
       // Intersect the in-bounds with the mask specified as an op parameter.

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
-/// Helper function that creates a memref::DimOp or tensor::DimOp depending on
-/// the type of `source`.
-Value mlir::vector::createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
-                                      int64_t dim) {
-  if (source.getType().isa<UnrankedMemRefType, MemRefType>())
-    return b.createOrFold<memref::DimOp>(loc, source, dim);
-  if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
-    return b.createOrFold<tensor::DimOp>(loc, source, dim);
-  llvm_unreachable("Expected MemRefType or TensorType");
-}
-
 /// Constructs a permutation map from memref indices to vector dimension.
 ///
 /// The implementation uses the knowledge of the mapping of enclosing loop to
 /// vector dimension. `enclosingLoopToVectorDim` carries this information as a
 /// map with:
 ///   - keys representing "vectorized enclosing loops";
 ///   - values representing the corresponding vector dimension.
 /// The algorithm traverses "vectorized enclosing loops" and extracts the

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   CastInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DataLayoutInterfaces.cpp
   DerivedAttributeOpInterface.cpp
   DestinationStyleOpInterface.cpp
   InferIntRangeInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   ParallelCombiningOpInterface.cpp
   RuntimeVerifiableOpInterface.cpp
   ShapedOpInterfaces.cpp
+  ShapedTypeInterfaces.cpp
   SideEffectInterfaces.cpp
   TilingInterface.cpp
   VectorInterfaces.cpp
   ViewLikeInterface.cpp
   )
 
 function(add_mlir_interface_library name)
   add_mlir_library(MLIR${name}

 add_mlir_interface_library(DerivedAttributeOpInterface)
 add_mlir_interface_library(DestinationStyleOpInterface)
 add_mlir_interface_library(InferIntRangeInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(LoopLikeInterface)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(RuntimeVerifiableOpInterface)
 add_mlir_interface_library(ShapedOpInterfaces)
+add_mlir_interface_library(ShapedTypeInterfaces)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(TilingInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 
 add_subdirectory(Utils)

mlir/lib/Interfaces/ShapedTypeInterfaces.cpp

+//===- ShapedTypeInterfaces.cpp - Interfaces for Shaped Types -------------===//
+//
+// 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/Interfaces/ShapedTypeInterfaces.h"
+
+#include "mlir/IR/Builders.h"
+
+#include "mlir/Interfaces/ShapedTypeInterfaces.cpp.inc"
+
+using namespace mlir;
+
+OpFoldResult mlir::reifyShapeDim(OpBuilder &builder, Location loc,
+                                 Value shapedValue, OpFoldResult dim) {
+  auto shapedTypeInterface =
+      shapedValue.getType().cast<ReifyShapeDimTypeInterface>();
+  return shapedTypeInterface.reifyShapeDim(builder, loc, shapedValue, dim);
+}
+
+OpFoldResult mlir::reifyShapeDim(OpBuilder &builder, Location loc,
+                                 Value shapedValue, int64_t dim) {
+  return reifyShapeDim(builder, loc, shapedValue, builder.getIndexAttr(dim));
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 td_library(
     name = "ShapedOpInterfacesTdFiles",
     srcs = ["include/mlir/Interfaces/ShapedOpInterfaces.td"],
     includes = ["include"],
     deps = [":OpBaseTdFiles"],
 )
 
 td_library(
+    name = "ShapedTypeInterfacesTdFiles",
+    srcs = ["include/mlir/Interfaces/ShapedTypeInterfaces.td"],
+    includes = ["include"],
+    deps = [":OpBaseTdFiles"],
+)
+
+td_library(
     name = "ParallelCombiningOpInterfaceTdFiles",
     srcs = ["include/mlir/Interfaces/ParallelCombiningOpInterface.td"],
     includes = ["include"],
     deps = [":OpBaseTdFiles"],
 )
 
 td_library(
     name = "RuntimeVerifiableOpInterfaceTdFiles",
deps = [
         ":LinalgDialect",
         ":LinalgTransforms",
         ":LinalgUtils",
         ":MathDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":SCFTransforms",
+        ":ShapedTypeInterfaces",
         ":SparseTensorDialect",
         ":SparseTensorEnums",
         ":SparseTensorPassIncGen",
         ":SparseTensorUtils",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorDialect",
cc_library(
     deps = [
         ":IR",
         ":ShapedOpInterfacesIncGen",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "ShapedTypeInterfaces",
+    srcs = ["lib/Interfaces/ShapedTypeInterfaces.cpp"],
+    hdrs = ["include/mlir/Interfaces/ShapedTypeInterfaces.h"],
+    includes = ["include"],
+    deps = [
+        ":IR",
+        ":ShapedTypeInterfacesIncGen",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
     name = "ParallelCombiningOpInterface",
     srcs = ["lib/Interfaces/ParallelCombiningOpInterface.cpp"],
     hdrs = ["include/mlir/Interfaces/ParallelCombiningOpInterface.h"],
     includes = ["include"],
     deps = [
         ":IR",
         ":ParallelCombiningOpInterfaceIncGen",
         "//llvm:Support",
deps = [
         ":BufferizationTransforms",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
+        ":ShapedTypeInterfaces",
         ":SideEffectInterfaces",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorDialect",
         ":VectorEnumsIncGen",
         ":VectorInterfaces",
         ":VectorPassIncGen",
gentbl_cc_library(
     td_file = "include/mlir/Dialect/Tensor/IR/TensorOps.td",
     deps = [":TensorOpsTdFiles"],
 )
 
 cc_library(
     name = "TensorDialect",
     srcs = [
         "include/mlir/Transforms/InliningUtils.h",
+        "lib/Dialect/Tensor/IR/ReifyShapeDimTypeInterfaceImpl.cpp",
         "lib/Dialect/Tensor/IR/TensorDialect.cpp",
         "lib/Dialect/Tensor/IR/TensorOps.cpp",
     ],
-    hdrs = ["include/mlir/Dialect/Tensor/IR/Tensor.h"],
+    hdrs = [
+        "include/mlir/Dialect/Tensor/IR/ReifyShapeDimTypeInterfaceImpl.h",
+        "include/mlir/Dialect/Tensor/IR/Tensor.h",
+    ],
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithDialect",
         ":ArithUtils",
         ":CastOpInterfaces",
         ":ComplexDialect",
         ":ControlFlowInterfaces",
         ":DestinationStyleOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":ParallelCombiningOpInterface",
         ":ShapedOpInterfaces",
+        ":ShapedTypeInterfaces",
         ":SideEffectInterfaces",
         ":Support",
         ":TensorOpsIncGen",
         ":TilingInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )
deps = [
         ":BufferizationTransforms",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
+        ":ShapedTypeInterfaces",
         ":TensorDialect",
         ":TensorPassIncGen",
         ":TilingInterface",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
tbl_outs = [
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Interfaces/ShapedOpInterfaces.td",
     deps = [":ShapedOpInterfacesTdFiles"],
 )
 
 gentbl_cc_library(
+    name = "ShapedTypeInterfacesIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["-gen-type-interface-decls"],
+            "include/mlir/Interfaces/ShapedTypeInterfaces.h.inc",
+        ),
+        (
+            ["-gen-type-interface-defs"],
+            "include/mlir/Interfaces/ShapedTypeInterfaces.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Interfaces/ShapedTypeInterfaces.td",
+    deps = [":ShapedTypeInterfacesTdFiles"],
+)
+
+gentbl_cc_library(
     name = "ParallelCombiningOpInterfaceIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             ["-gen-op-interface-decls"],
             "include/mlir/Interfaces/ParallelCombiningOpInterface.h.inc",
         ),
         (
deps = [
         ":LinalgInterfacesIncGen",
         ":LinalgNamedStructuredOpsYamlIncGen",
         ":LinalgOpsIncGen",
         ":LinalgStructuredOpsIncGen",
         ":MathDialect",
         ":MemRefDialect",
         ":Parser",
         ":SCFDialect",
+        ":ShapedTypeInterfaces",
         ":SideEffectInterfaces",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         ":TilingInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
deps = [
         ":ArithUtils",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
+        ":ShapedTypeInterfaces",
         ":TensorDialect",
         ":TensorUtils",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "LinalgTransforms",
deps = [
         ":MaskableOpInterface",
         ":MathDialect",
         ":MemRefDialect",
         ":MemRefTransforms",
         ":Pass",
         ":SCFDialect",
         ":SCFTransforms",
         ":SCFUtils",
+        ":ShapedTypeInterfaces",
         ":SparseTensorDialect",
         ":Support",
         ":TensorDialect",
         ":TensorTilingInterfaceImpl",
         ":TensorTransforms",
         ":TensorUtils",
         ":TilingInterface",
         ":TransformUtils",
deps = [
         ":ConversionPassIncGen",
         ":FuncDialect",
         ":FuncToLLVM",
         ":IR",
         ":LLVMDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
+        ":ShapedTypeInterfaces",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorDialect",
         ":VectorTransforms",
         "//llvm:Core",
         "//llvm:Support",
     ],
srcs = glob(
         [
             "lib/Dialect/MemRef/IR/*.cpp",
             "lib/Dialect/MemRef/IR/*.h",
             "lib/Dialect/MemRef/Utils/*.cpp",
         ],
     ),
     hdrs = [
         "include/mlir/Dialect/MemRef/IR/MemRef.h",
+        "include/mlir/Dialect/MemRef/IR/ReifyShapeDimTypeInterfaceImpl.h",
         "include/mlir/Dialect/MemRef/Utils/MemRefUtils.h",
     ],
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":ArithUtils",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefBaseIncGen",
         ":MemRefOpsIncGen",
         ":ShapedOpInterfaces",
+        ":ShapedTypeInterfaces",
         ":ViewLikeInterface",
         "//llvm:Support",
         "//llvm:TargetParser",
     ],
 )
 
 gentbl_cc_library(
     name = "MemRefPassIncGen",