Fix tests

Author: nicolasvasilache

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

@@ -214,9 +214,10 @@ def Vector_ContractionOp :
     // IndexingMapOpInterface interface methods implementation.
     //===------------------------------------------------------------------===//
     ArrayRef<int64_t> getShape(OpOperand * opOperand) {
-      assert(opOperand->getOwner() == this->getOperation());
       Type t = opOperand->get().getType();
-      return cast<VectorType>(t).getShape();
+      if (auto vt = dyn_cast<VectorType>(t))
+        return vt.getShape();
+      return {};
     }
   }];
 

mlir/include/mlir/Interfaces/IndexingMapOpInterface.h

@@ -14,6 +14,13 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 
+namespace mlir {
+namespace detail {
+/// Verify that `op` conforms to the invariants of StructuredOpInterface
+LogicalResult verifyIndexingMapOpInterface(Operation *op);
+} // namespace detail
+} // namespace mlir
+
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/IndexingMapOpInterface.h.inc"
 

mlir/include/mlir/Interfaces/IndexingMapOpInterface.td

@@ -83,23 +83,6 @@ def IndexingMapOpInterface : OpInterface<"IndexingMapOpInterface"> {
         return concatAffineMaps(maps, $_op.getContext());
       }]
     >,
-    InterfaceMethod<
-      /*desc=*/[{
-        Like `getShape`, but only returns statically-known information, without
-        generating any new IR. For each shape dimension, returns >=0 if that
-        dimension is statically known, or ShapedType::kDynamic otherwise.
-      }],
-      /*retTy=*/"SmallVector<int64_t>",
-      /*methodName=*/"getStaticShape",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        SmallVector<int64_t> res;
-        for (OpOperand &opOperand : this->getOperation()->getOpOperands())
-          llvm::append_range(res, $_op.getShape(&opOperand));
-        return res;
-      }]
-    >,
     InterfaceMethod<
       /*desc=*/[{
         Hook to provide a custom AffineMap used to construct the
@@ -126,23 +109,46 @@ def IndexingMapOpInterface : OpInterface<"IndexingMapOpInterface"> {
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Returns the statically-known loop ranges. Composes
-        `getShapesToLoopsMap()` with the result of `getStaticShape`.
+        Returns the static shape of the underlying operand (note this is
+        op-specific behavior).
+        Returns ShapedType::kDynamic for non-statically-known loop ranges.
+      }],
+      /*retTy=*/"SmallVector<int64_t>",
+      /*methodName=*/"getStaticOperandShape",
+      /*args=*/(ins "OpOperand*":$opOperand),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        SmallVector<int64_t> res;
+        llvm::append_range(res, $_op.getShape(opOperand));
+        return res;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns loop ranges by composing `getShapesToLoopsMap()` with the
+        flattened list of operand shapes.
         Returns ShapedType::kDynamic for non-statically-known loop ranges.
-        This is expected to be called by a valid Linalg op
       }],
-      /*retTy=*/"SmallVector<int64_t, 4>",
+      /*retTy=*/"SmallVector<int64_t>",
       /*methodName=*/"getStaticLoopRanges",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        SmallVector<int64_t> viewSizes = $_op.getStaticShape();
+        SmallVector<int64_t> allShapesSizes;
+        for (OpOperand &opOperand : this->getOperation()->getOpOperands())
+          llvm::append_range(allShapesSizes, $_op.getShape(&opOperand));
         AffineMap invertedMap = $_op.getShapesToLoopsMap();
-        assert(invertedMap && "expected a valid Linalg op to call the method");
-        return invertedMap.compose(viewSizes);
+        assert(invertedMap && "expected a valid op");
+        return invertedMap.compose(allShapesSizes);
       }]
-    >,
+    >
   ];
+  let extraClassDeclaration = [{
+    // Verifier implementation for IndexingMapOpInterface.
+    // This must be called manually as part of other verifiers so that the
+    // verification order, and meaningful error messages, are not preempted.
+    LogicalResult verifyImpl();
+  }];
 }
 
-#endif // MLIR_INTERFACES_INDEXING_MAP_OP_INTERFACE
+#endif // MLIR_INTERFACES_INDEXING_MAP_OP_INTERFACE

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

@@ -1215,106 +1215,27 @@ LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
     if (failed(linalgOp.verifyIndexingMapRequiredAttributes()))
       return failure();
 
-  // All input/output operands must be indexed.
-  if (static_cast<int64_t>(linalgOp.getIndexingMapsArray().size()) !=
-      linalgOp->getNumOperands())
-    return op->emitOpError("expected the number of indexing_map (")
-           << linalgOp.getIndexingMapsArray().size()
-           << ") to be equal to the number of input/output operands ("
-           << linalgOp->getNumOperands() << ")";
+  // Delayed calling of IndexingMapOpInterface::verifyImpl.
+  if (failed(cast<IndexingMapOpInterface>(op).verifyImpl()))
+    return failure();
 
   // Set this flag if this op has user defined maps. This is required to guard
   // the below error condition which assume default indexing maps.
   for (OpOperand &opOperand : linalgOp->getOpOperands()) {
     AffineMap indexingMap = linalgOp.getMatchingIndexingMap(&opOperand);
-
-    // Symbols disallowed.
-    if (indexingMap.getNumSymbols() != 0)
-      return op->emitOpError("unexpected symbols in indexing_map #")
-             << opOperand.getOperandNumber();
-
     // Domain must be consistent.
     unsigned numLoops = linalgOp.getNumLoops();
     if (indexingMap.getNumDims() != numLoops)
       return op->emitOpError("expected indexing_map #")
              << opOperand.getOperandNumber() << " to have " << numLoops
              << " dim(s) to match the number of loops";
-
-    int64_t rank = linalgOp.getRank(&opOperand);
-
-    if (indexingMap.getNumResults() != rank)
-      return op->emitOpError("expected operand rank (")
-             << rank << ") to match the result rank of indexing_map #"
-             << opOperand.getOperandNumber() << " ("
-             << indexingMap.getNumResults() << ")";
   }
   SmallVector<unsigned> redDims;
   linalgOp.getReductionDims(redDims);
 
   if (!linalgOp.getShapesToLoopsMap())
     return op->emitOpError("expected the shape-to-loops map to be non-null");
 
-  // Check if given shapes match to inferred shapes.
-  SmallVector<int64_t, 4> endLoopRangeValues = linalgOp.getStaticLoopRanges();
-  SmallVector<int64_t, 4> startLoopRangeValues(endLoopRangeValues.size(), 0);
-  // Verify only static cases since we can't get exact dimension sizes and
-  // loop ranges for dynamic cases in this stage.
-  if (llvm::none_of(endLoopRangeValues, ShapedType::isDynamic)) {
-    for (int64_t &range : endLoopRangeValues)
-      range -= 1;
-    for (OpOperand &opOperand : linalgOp->getOpOperands()) {
-      AffineMap indexingMap = linalgOp.getMatchingIndexingMap(&opOperand);
-      SmallVector<int64_t, 4> startIndices =
-          indexingMap.compose(startLoopRangeValues);
-      SmallVector<int64_t, 4> endIndices =
-          indexingMap.compose(endLoopRangeValues);
-      ArrayRef<int64_t> shape = linalgOp.getShape(&opOperand);
-      for (auto dim : llvm::seq<int64_t>(0, shape.size())) {
-        // Ignore dynamic dimension or the case that the dimension size is 0
-        if (ShapedType::isDynamic(shape[dim]) || shape[dim] == 0)
-          continue;
-
-        // The first index or last index should be the maximum or the minimum in
-        // the inferred index ranges since the range is increasing or
-        // decreasing. The size of dimensions of input/output operands and the
-        // maximum value + 1 in the inferred range should be the same. But, for
-        // now we check if the inferred ranges are in boundary of input/output
-        // operands' size or not in case that Affine Expressions are complicated
-        // such as d0 * 3
-        // + d1 since it is not easy to handle the issues.
-        // Found the case that this solution can't check, for example, (d0, d1)
-        // -> (d1 - d0)
-        int64_t inferredDimSize =
-            std::max(startIndices[dim], endIndices[dim]) + 1;
-        if (std::min(startIndices[dim], endIndices[dim]) < 0) {
-          std::string mapStr;
-          {
-            llvm::raw_string_ostream os(mapStr);
-            os << indexingMap;
-          }
-          return op->emitOpError(
-                     "unexpected result less than 0 at expression #")
-                 << dim << " in " << mapStr;
-        }
-        if (isa<AffineDimExpr>(indexingMap.getResult(dim))) {
-          if (inferredDimSize != shape[dim]) {
-            return op->emitOpError("inferred input/output operand #")
-                   << opOperand.getOperandNumber() << " has shape's dimension #"
-                   << dim << " to be " << inferredDimSize << ", but found "
-                   << shape[dim];
-          }
-        } else {
-          if (inferredDimSize > shape[dim]) {
-            return op->emitOpError("inferred input/output operand #")
-                   << opOperand.getOperandNumber() << " has shape's dimension #"
-                   << dim << " to be greater than or equal to "
-                   << inferredDimSize << ", but found " << shape[dim];
-          }
-        }
-      }
-    }
-  }
-
   // Check the region has exactly one block.
   if (linalgOp->getNumRegions() != 1 ||
       !llvm::hasSingleElement(linalgOp->getRegion(0)))

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

@@ -398,7 +398,10 @@ linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
     return rewriter.notifyMatchFailure(genericOp,
                                        "invalid indexing maps for operation");
   }
-  SmallVector<int64_t> dims = genericOp.getStaticShape();
+
+  SmallVector<int64_t> allShapesSizes;
+  for (OpOperand &opOperand : genericOp->getOpOperands())
+    llvm::append_range(allShapesSizes, genericOp.getShape(&opOperand));
 
   // 1a. Get the allowed list of dimensions to drop from the `options`.
   SmallVector<unsigned> allowedUnitDims = options.controlFn(genericOp);
@@ -411,7 +414,7 @@ linalg::dropUnitDims(RewriterBase &rewriter, GenericOp genericOp,
   llvm::SmallDenseSet<unsigned> unitDims;
   for (const auto &expr : enumerate(invertedMap.getResults())) {
     if (AffineDimExpr dimExpr = dyn_cast<AffineDimExpr>(expr.value())) {
-      if (dims[dimExpr.getPosition()] == 1 &&
+      if (allShapesSizes[dimExpr.getPosition()] == 1 &&
           unitDimsFilter.count(expr.index()))
         unitDims.insert(expr.index());
     }

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

@@ -31,6 +31,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
@@ -2217,7 +2218,9 @@ static LogicalResult vectorizeLinalgOpPrecondition(
     LinalgOp linalgOp, ArrayRef<int64_t> inputVectorSizes,
     bool vectorizeNDExtract, bool flatten1DDepthwiseConv) {
   // tensor with dimension of 0 cannot be vectorized.
-  if (llvm::is_contained(linalgOp.getStaticShape(), 0))
+  if (llvm::any_of(linalgOp->getOpOperands(), [&](OpOperand &operand) {
+        return llvm::is_contained(linalgOp.getShape(&operand), 0);
+      }))
     return failure();
   // Check API contract for input vector sizes.
   if (!inputVectorSizes.empty() &&

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -1063,7 +1063,8 @@ LogicalResult ContractionOp::verify() {
   if (!isSupportedCombiningKind(getKind(), elementType))
     return emitOpError("unsupported contraction type");
 
-  return success();
+  // Delayed calling of IndexingMapOpInterface::verifyImpl.
+  return cast<IndexingMapOpInterface>(this->getOperation()).verifyImpl();
 }
 
 // MaskableOpInterface methods.

mlir/lib/Interfaces/IndexingMapOpInterface.cpp

@@ -13,3 +13,113 @@ using namespace mlir;
 namespace mlir {
 #include "mlir/Interfaces/IndexingMapOpInterface.cpp.inc"
 } // namespace mlir
+
+LogicalResult mlir::IndexingMapOpInterface::verifyImpl() {
+  // All input/output operands must be indexed.
+  if (static_cast<int64_t>(getIndexingMapsArray().size()) !=
+      getOperation()->getNumOperands())
+    return this->emitOpError("expected the number of indexing_map (")
+           << getIndexingMapsArray().size()
+           << ") to be equal to the number of input/output operands ("
+           << getOperation()->getNumOperands() << ")";
+
+  AffineMap invertedMap = getShapesToLoopsMap();
+  if (!invertedMap) {
+    std::string str;
+    llvm::raw_string_ostream os(str);
+    getLoopsToShapesMap().print(os);
+    return this->emitOpError("invalid indexing maps are non-invertible: ")
+           << "(" << str << ")";
+  }
+
+  SmallVector<int64_t> endLoopRangeValues = getStaticLoopRanges();
+
+  // Set this flag if this op has user defined maps. This is required to guard
+  // the below error condition which assume default indexing maps.
+  for (OpOperand &opOperand : getOperation()->getOpOperands()) {
+    AffineMap indexingMap = getMatchingIndexingMap(&opOperand);
+
+    // Symbols disallowed.
+    if (indexingMap.getNumSymbols() != 0)
+      return getOperation()->emitOpError("unexpected symbols in indexing_map #")
+             << opOperand.getOperandNumber();
+
+    // Domain must be consistent.
+    if (indexingMap.getNumDims() != endLoopRangeValues.size())
+      return getOperation()->emitOpError("expected indexing_map #")
+             << opOperand.getOperandNumber() << " to have "
+             << endLoopRangeValues.size()
+             << " dim(s) to match the number of loops";
+
+    SmallVector<int64_t> shape = getStaticOperandShape(&opOperand);
+    int64_t rank = shape.size();
+
+    if (indexingMap.getNumResults() != rank)
+      return getOperation()->emitOpError("expected operand rank (")
+             << rank << ") to match the result rank of indexing_map #"
+             << opOperand.getOperandNumber() << " ("
+             << indexingMap.getNumResults() << ")";
+  }
+
+  // Check if given shapes match to inferred shapes.
+  SmallVector<int64_t> startLoopRangeValues(endLoopRangeValues.size(), 0);
+  // Verify only static cases since we can't get exact dimension sizes and
+  // loop ranges for dynamic cases in this stage.
+  if (llvm::none_of(endLoopRangeValues, ShapedType::isDynamic)) {
+    // Exclusive end range.
+    for (int64_t &range : endLoopRangeValues)
+      range -= 1;
+    for (OpOperand &opOperand : getOperation()->getOpOperands()) {
+      AffineMap indexingMap = getMatchingIndexingMap(&opOperand);
+      SmallVector<int64_t> startIndices =
+          indexingMap.compose(startLoopRangeValues);
+      SmallVector<int64_t> endIndices = indexingMap.compose(endLoopRangeValues);
+      SmallVector<int64_t> shape = getStaticOperandShape(&opOperand);
+      for (auto dim : llvm::seq<int64_t>(0, shape.size())) {
+        // Ignore dynamic dimension or the case that the dimension size is 0
+        if (ShapedType::isDynamic(shape[dim]) || shape[dim] == 0)
+          continue;
+
+        // The first index or last index should be the maximum or the minimum in
+        // the inferred index ranges since the range is increasing or
+        // decreasing. The size of dimensions of input/output operands and the
+        // maximum value + 1 in the inferred range should be the same. But, for
+        // now we check if the inferred ranges are in boundary of input/output
+        // operands' size or not in case that Affine Expressions are complicated
+        // such as d0 * 3
+        // + d1 since it is not easy to handle the issues.
+        // Found the case that this solution can't check, for example, (d0, d1)
+        // -> (d1 - d0)
+        int64_t inferredDimSize =
+            std::max(startIndices[dim], endIndices[dim]) + 1;
+        if (std::min(startIndices[dim], endIndices[dim]) < 0) {
+          std::string mapStr;
+          {
+            llvm::raw_string_ostream os(mapStr);
+            os << indexingMap;
+          }
+          return this->emitOpError(
+                     "unexpected result less than 0 at expression #")
+                 << dim << " in " << mapStr;
+        }
+        if (isa<AffineDimExpr>(indexingMap.getResult(dim))) {
+          if (inferredDimSize != shape[dim]) {
+            return this->emitOpError("inferred input/output operand #")
+                   << opOperand.getOperandNumber() << " has shape's dimension #"
+                   << dim << " to be " << inferredDimSize << ", but found "
+                   << shape[dim];
+          }
+        } else {
+          if (inferredDimSize > shape[dim]) {
+            return this->emitOpError("inferred input/output operand #")
+                   << opOperand.getOperandNumber() << " has shape's dimension #"
+                   << dim << " to be greater than or equal to "
+                   << inferredDimSize << ", but found " << shape[dim];
+          }
+        }
+      }
+    }
+  }
+
+  return success();
+}

mlir/test/Dialect/Linalg/invalid.mlir

@@ -151,7 +151,7 @@ func.func @generic_result_0_element_type(%arg0: memref<?xf32, affine_map<(i)[off
 // -----
 
 func.func @generic_singular_maps(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>, %arg1: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
-  // expected-error @+1 {{expected the shape-to-loops map to be non-null}}
+  // expected-error @+1 {{invalid indexing maps are non-invertible: ((d0, d1) -> (d0 + d1, d0 + d1))}}
   linalg.generic {
     indexing_maps =  [
       affine_map<(i, j) -> (i + j)>,