[Mlir-commits] [mlir] [MLIR][IndexingMapOpInterface]: Validate maps and operands before composing loop ranges (PR #173434)

[lonely eagle]

================
@@ -23,44 +23,44 @@ LogicalResult mlir::IndexingMapOpInterface::verifyImpl() {
            << ") 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();
+  SmallVector<int64_t> allShapesSizes;
 
-  // 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);
+    SmallVector<int64_t> shape = getStaticOperandShape(&opOperand);
+    int64_t rank = shape.size();
 
     // Symbols disallowed.
     if (indexingMap.getNumSymbols() != 0)
-      return getOperation()->emitOpError("unexpected symbols in indexing_map #")
+      return this->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();
-
+    // Result rank must match operand rank.
     if (indexingMap.getNumResults() != rank)
-      return getOperation()->emitOpError("expected operand rank (")
-             << rank << ") to match the result rank of indexing_map #"
-             << opOperand.getOperandNumber() << " ("
+      return this->emitOpError("expected operand #")
+             << opOperand.getOperandNumber() << " rank (" << rank
+             << ") to match the result rank of indexing_map ("
              << indexingMap.getNumResults() << ")";
+
+    llvm::append_range(allShapesSizes, shape);
   }
 
+  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 = invertedMap.compose(allShapesSizes);
+
+  if (invertedMap.getNumResults() != endLoopRangeValues.size())
----------------
linuxlonelyeagle wrote:

Shouldn't the condition we are checking here be: `indexingMap.getNumDims() != endLoopRangeValues.size()`? Since we checked the size of the result in the first half, we should check the size of the dimensions (dims) in the second half.

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