[MLIR] Avoid creation of buggy affine maps when incorrect values of number of dimensions and number of symbols are provided.

[arnab-oss]

Summary

We check whether the maximum index of dimensional identifier present
in the result expressions is less than dimCount (number of dimensional
identifiers) argument passed in the AffineMap::get() and the maximum index
of symbolic identifier present in the result expressions is less than
symbolCount (number of symbolic identifiers) argument passed in AffineMap::get().

[rriddle]

mlir/lib/IR/MLIRContext.cpp

1021–1040

Instead of recursion, can you use a worklist?

[bondhugula]

mlir/lib/IR/MLIRContext.cpp

1021–1040

The recursion here is no different than the recursion that happens when parsing affine expressions, walking, simplifying, or processing affine expressions. Using a worklist would be harder to read. Note that this is recursing over existing IR (left and right sub-trees of binary exprs) -- it's bounded and O(IR size). Such recursion is allowed, right?

[bondhugula]

mlir/lib/IR/MLIRContext.cpp

39

I think you won't need this include.

[nicolasvasilache]

mlir/lib/IR/MLIRContext.cpp

1014

You can write this check in 3 lines by calling AffineMap::inferFromExprList and checking against the returned AffineMap.

[arnab-oss]

Addressed comments.

[arnab-oss]

Removed unnecessary #include<algorithm>.

[arnab-oss]

ping @rriddle and @nicolasvasilache can you please take a look at the replies of your respective comments.

mlir/lib/IR/MLIRContext.cpp

1014

AffineMap::inferFromExprList calls AffineMap::get() internally, which causes a layering issue as we are indirectly calling AffineMap::get() from inside AffineMap::get(). But before calling AffineMap::get() using a similar traversal. So, first creating a correct map separately, and then checking whether numDims and numSymbols follow the constraints will do an extra map allocation, as it will call the getImpl() for creating the correct map. Our aim is to not create any map at all if the arguments are wrong.
CC - @bondhugula

[nicolasvasilache]

mlir/lib/IR/MLIRContext.cpp

1014

Sure, but I imagine you can think of ways of layering this differently?

Absolute worst case, you can expose:

template <typename AffineExprContainer>
static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
                               int64_t &maxDim, int64_t &maxSym) {

as a util and reuse it rather than reimplement a 5 line iterative algorithm with a longer recursive one.

[arnab-oss]

Addressed comments by @nicolasvasilache and @bondhugula

[arnab-oss]

@nicolasvasilache, @bondhugula and @rriddle, I have addressed all your comments. Please take a look.

[nicolasvasilache]

thanks!

[arnab-oss]

Rebased on latest master.

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

550 ↗

(On Diff #389883)

Missing doc comment here - it wasn't there earlier, but we should add one now that this is publicly exposed.

[arnab-oss]

Added doc comments for the getMaxDimAndSymbol().

[bondhugula]

mlir/include/mlir/IR/AffineMap.h

549 ↗

(On Diff #390098)

Nit: position -> positions

in the `exprsLists` -> in `exprsLists`

[mehdi_amini]

mlir/lib/IR/MLIRContext.cpp

1014

"nodiscard" isn't universally supported and it is a c++17 extension, so someone removed it after you committed this.

But more importantly: this is not desirable! What it means is to force the compiler to emit this function in the binary even if unused, there is a rarely a reason to do that, it'll just bloat the binary size.

I assume that the problem is the warning about this function being unused, in this case you can:

1. Surround the function between #ifndef NDEBUG / #endif
2. Use the LLVM_ATTRIBUTE_UNUSED decorator in front the function.

mlir/lib/IR/MLIRContext.cpp

 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/RWMutex.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/raw_ostream.h"
 #include <memory>
 
 #define DEBUG_TYPE "mlircontext"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 using llvm::hash_combine;
 using llvm::hash_combine_range;
return safeGetOrCreate(
 
         // Initialize the memory using placement new.
         new (res)
             detail::AffineMapStorage{dimCount, symbolCount, results, context};
         return AffineMap(res);
       });
 }
 
+/// Check whether the arguments passed to the AffineMap::get() are consistent.
+/// This method checks whether the highest index of dimensional identifier
+/// present in result expressions is less than `dimCount` and the highest index
+/// of symbolic identifier present in result expressions is less than
+/// `symbolCount`.
+[[nodiscard]] static bool willBeValidAffineMap(unsigned dimCount,
+                                               unsigned symbolCount,
+                                               ArrayRef<AffineExpr> results) {
+  unsigned numDimensions = 0;
+  unsigned numSymbols = 0;
+  std::function<void(AffineExpr)> calculateDimAndSymbolCount =
+      [&](AffineExpr expr) {
+        if (auto binOpExpr = expr.dyn_cast<AffineBinaryOpExpr>()) {
+          calculateDimAndSymbolCount(binOpExpr.getLHS());
+          calculateDimAndSymbolCount(binOpExpr.getRHS());
+          return;
+        }
+        if (auto dimExpr = expr.dyn_cast<AffineDimExpr>())
+          numDimensions = std::max(numDimensions, 1 + dimExpr.getPosition());
+        if (auto symbolExpr = expr.dyn_cast<AffineSymbolExpr>()) {
+          numSymbols = std::max(numSymbols, 1 + symbolExpr.getPosition());
+        }
+      };
+
+  for (AffineExpr expr : results)
+    calculateDimAndSymbolCount(expr);
+
+  if ((numDimensions > dimCount) || (numSymbols > symbolCount)) {
+    LLVM_DEBUG(
+        llvm::dbgs()
+        << "maximum dimensional identifier position in result expression must "
+           "be less than `dimCount` and maximum symbolic identifier position "
+           "in result expression must be less than `symbolCount`\n");
+    return false;
+  }
+  return true;
+}
+
 AffineMap AffineMap::get(MLIRContext *context) {
   return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context);
 }
 
 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
                          MLIRContext *context) {
   return getImpl(dimCount, symbolCount, /*results=*/{}, context);
 }
 
 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
                          AffineExpr result) {
+  assert(willBeValidAffineMap(dimCount, symbolCount, {result}));
   return getImpl(dimCount, symbolCount, {result}, result.getContext());
 }
 
 AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
                          ArrayRef<AffineExpr> results, MLIRContext *context) {
+  assert(willBeValidAffineMap(dimCount, symbolCount, results));
   return getImpl(dimCount, symbolCount, results, context);
 }
 
 //===----------------------------------------------------------------------===//
 // Integer Sets: these are allocated into the bump pointer, and are immutable.
 // Unlike AffineMap's, these are uniqued only if they are small.
 //===----------------------------------------------------------------------===//