[mlir][linalg] Add isPermutation helper (NFC).

[gysit]

Summary

Add a helper method to check if an index vector contains a permutation of its indices. Additionally, refactor applyPermutationToVector to take int64_t.

[gysit]

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

34

AFAIK the context is needed to create an AffineMap. Optionally, we may also add an assertion inside of applyPermutationToVector, but this would also require adding context to the argument list.

[nicolasvasilache]

you could also just traverse by recording the min, max and unique counts to avoid creating the AffineMap in the first place.

This works too.

[gysit]

Verify permutation without using AffineMap.

[gysit]

Add unsigned to signed casts.

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

 
 namespace linalg {
 class LinalgDependenceGraph;
 
 //===----------------------------------------------------------------------===//
 // General utilities
 //===----------------------------------------------------------------------===//
 
+/// Check if `permutation` is a permutation of the range
+/// `[0, permutation.size())`.
+bool isPermutation(ArrayRef<int64_t> permutation);
+
 /// Apply the permutation defined by `permutation` to `inVec`.
 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
 /// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.
 template <typename T, unsigned N>
 void applyPermutationToVector(SmallVector<T, N> &inVec,
-                              ArrayRef<unsigned> permutation) {
+                              ArrayRef<int64_t> permutation) {
   SmallVector<T, N> auxVec(inVec.size());
-  for (unsigned i = 0; i < permutation.size(); ++i)
-    auxVec[i] = inVec[permutation[i]];
+  for (auto en : enumerate(permutation))
+    auxVec[en.index()] = inVec[en.value()];
   inVec = auxVec;
 }
 
 /// 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);
 
 /// Given an operation, retrieves the value of each dynamic dimension through

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

 //===----------------------------------------------------------------------===//
 
 FailureOr<TileLoopNest>
 mlir::linalg::tileConsumerAndFuseProducers(OpBuilder &b, LinalgOp consumerOp,
                                            ArrayRef<int64_t> tileSizes,
                                            ArrayRef<int64_t> tileInterchange) {
   assert(tileSizes.size() == tileInterchange.size() &&
          "expect the number of tile sizes and interchange dims to match");
+  assert(isPermutation(tileInterchange) &&
+         "expect tile interchange is a permutation");
 
   // Create an empty tile loop nest.
   TileLoopNest tileLoopNest(consumerOp);
 
   // Search the number of outer parallel loops to separate them from possible
   // inner reduction dimensions.
   SmallVector<StringAttr> iterTypes =
       llvm::to_vector<6>(consumerOp.iterator_types().getAsRange<StringAttr>());
-  applyPermutationToVector(
-      iterTypes,
-      SmallVector<unsigned>(tileInterchange.begin(), tileInterchange.end()));
+  applyPermutationToVector(iterTypes, tileInterchange);
   auto *it = find_if(iterTypes, [&](StringAttr iterType) {
     return !isParallelIterator(iterType);
   });
   int64_t split = std::distance(iterTypes.begin(), it);
 
   // Tile the outer parallel loops and fuse the output operands.
   SmallVector<int64_t> outerTileSizes;
   outerTileSizes.append(tileSizes.begin(), tileSizes.begin() + split);
SmallVector<int64_t> rootTileSizes(
         tileSizes.begin(), tileSizes.begin() + rootOp.getNumLoops());
     SmallVector<int64_t> rootInterchange =
         tileInterchange.empty()
             ? llvm::to_vector<6>(llvm::seq<int64_t>(0, tileSizes.size()))
             : SmallVector<int64_t>(tileInterchange.begin(),
                                    tileInterchange.begin() +
                                        rootOp.getNumLoops());
 
-    // As a tiling can only tile a loop dimension once, `rootInterchange` has to
-    // be a permutation of the `rootOp` loop dimensions.
-    SmallVector<AffineExpr> rootInterchangeExprs;
-    transform(rootInterchange, std::back_inserter(rootInterchangeExprs),
-              [&](int64_t dim) { return b.getAffineDimExpr(dim); });
-    AffineMap rootInterchangeMap = AffineMap::get(
-        rootOp.getNumLoops(), 0, rootInterchangeExprs, funcOp.getContext());
-    if (!rootInterchangeMap.isPermutation())
+    // Check `rootInterchange` is a permutation of the `rootOp` loop dimensions.
+    // It has to be a permutation since the tiling cannot tile the same loop
+    // dimension multiple times.
+    if (!isPermutation(rootInterchange))
       return notifyFailure(
           "expect the tile interchange permutes the root loops");
 
     // Tile `rootOp` and fuse its producers.
     if (failed(tileConsumerAndFuseProducers(b, rootOp, rootTileSizes,
                                             rootInterchange)))
       return notifyFailure("tileConsumerAndFuseProducers failed unexpectedly");
   }
 };
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTileAndFuseTensorOpsPass() {
   return std::make_unique<LinalgTileAndFuseTensorOps>();
 }

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

void mlir::linalg::interchangeGenericOp(PatternRewriter &rewriter,
   }
   genericOp->setAttr(getIndexingMapsAttrName(),
                      ArrayAttr::get(context, newIndexingMaps));
 
   // 3. Compute the interchanged iterator types.
   ArrayRef<Attribute> itTypes = genericOp.iterator_types().getValue();
   SmallVector<Attribute, 4> itTypesVector;
   llvm::append_range(itTypesVector, itTypes);
-  applyPermutationToVector(itTypesVector, interchangeVector);
+  SmallVector<int64_t> permutation(interchangeVector.begin(),
+                                   interchangeVector.end());
+  applyPermutationToVector(itTypesVector, permutation);
   genericOp->setAttr(getIteratorTypesAttrName(),
                      ArrayAttr::get(context, itTypesVector));
 
   // 4. Transform the index operations by applying the permutation map.
   if (genericOp.hasIndexSemantics()) {
     // TODO: Remove the assertion and add a getBody() method to LinalgOp
     // interface once every LinalgOp has a body.
     assert(genericOp->getNumRegions() == 1 &&

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

for (auto pos : options.interchangeVector) {
         continue;
       interchangeVector.push_back(it->second);
     }
     // Interchange vector is guaranteed to be a permutation,
     // `inversePermutation` must succeed.
     invPermutationMap = inversePermutation(
         AffineMap::getPermutationMap(interchangeVector, b.getContext()));
     assert(invPermutationMap);
-    applyPermutationToVector(loopRanges, interchangeVector);
-    applyPermutationToVector(iteratorTypes, interchangeVector);
+    SmallVector<int64_t> permutation(interchangeVector.begin(),
+                                     interchangeVector.end());
+    applyPermutationToVector(loopRanges, permutation);
+    applyPermutationToVector(iteratorTypes, permutation);
   }
 
   // 2. Create the tiled loops.
   LinalgOp res = op;
   SmallVector<Value, 4> ivs, tensorResults;
   auto tiledLoopBodyBuilder =
       [&](OpBuilder &b, Location loc, ValueRange localIvs,
           ValueRange operandValuesToUse) -> scf::ValueVector {

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

for (Range range : ranges) {
     ubs.emplace_back(range.size);
     steps.emplace_back(range.stride);
   }
 }
 
 namespace mlir {
 namespace linalg {
 
+bool isPermutation(ArrayRef<int64_t> permutation) {
+  // Count the number of appearances for all indices.
+  SmallVector<int64_t> indexCounts(permutation.size(), 0);
+  for (auto index : permutation) {
+    // Exit if the index is out-of-range.
+    if (index < 0 || index >= static_cast<int64_t>(permutation.size()))
+      return false;
+    indexCounts[index]++;
+  }
+  // Return true if all indices appear once.
+  return count(indexCounts, 1) == static_cast<int64_t>(permutation.size());
+}
+
 /// 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");