[mlir][Linalg] Handle null affine map returns from inversePermutation.

[mravishankar]

Summary

The inversePermutation method returns a null map on failure. Update
uses of this method within Linalg to handle this. In LinalgToLoops the
null return value was used to emit scalar code. Modify that to return
failure, and emit scalar implementation when affine map is "empty",
i.e. 1 dims, 0 symbols and no result exprs.
Also add test case for generating scalar code.

Depends On D77678

[mravishankar]

Rebase

[mehdi_amini]

Can you clarify the NFC here? Seems like you supply even a test to observe some behavior change?

[mravishankar]

Drop NFC and add additional checks at all uses of inversePermutation
to handle null return value.

[mravishankar]

rebase

[mravishankar]

Rebase

[mravishankar]

Rebase

[mravishankar]

Rebase

mlir/include/mlir/IR/AffineMap.h

 AffineMap simplifyAffineMap(AffineMap map);
 
 /// Returns a map with the same dimension and symbol count as `map`, but whose
 /// results are the unique affine expressions of `map`.
 AffineMap removeDuplicateExprs(AffineMap map);
 
 /// Returns a map of codomain to domain dimensions such that the first codomain
 /// dimension for a particular domain dimension is selected.
-/// Returns an empty map if the input map is empty or if `map` is not invertible
-/// (i.e. `map` does not contain a subset that is a permutation of full domain
-/// rank).
+/// Returns an empty map if the input map is empty.
+/// Returns null map (not empty map) if `map` is not invertible (i.e. `map` does
+/// not contain a subset that is a permutation of full domain rank).
 ///
 /// Prerequisites:
 ///   1. `map` has no symbols.
 ///
 /// Example 1:
 ///
 /// ```mlir
 ///    (d0, d1, d2) -> (d1, d1, d0, d2, d1, d2, d1, d0)

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

if (!areTensorOpsFusible(producer, consumer, consumerIdx))
     return {};
 
   MLIRContext *context = b.getContext();
   auto producerOp = cast<linalg::GenericOp>(producer.getOperation());
   auto consumerOp = cast<linalg::GenericOp>(consumer.getOperation());
   AffineMap consumerIndexMap = consumerOp.getIndexingMap(consumerIdx);
   AffineMap invProducerResultIndexMap =
       inversePermutation(producerOp.getOutputIndexingMap(0));
+  if (!invProducerResultIndexMap)
+    return {};
 
   // Compute the fused op operandslist by replacing the operand corresponding to
   // the result of the producer, with the operands of the producer.
   unsigned fusedArgsIn =
       producerOp.getNumInputs() + consumerOp.getNumInputs() - 1;
   auto fusedArgsOut = consumerOp.getNumOutputs();
   SmallVector<Value, 2> fusedOperandsList(consumerOp.getOperands());
   fusedOperandsList.erase(std::next(fusedOperandsList.begin(), consumerIdx));
for (auto operand : llvm::enumerate(op.getOperation()->getOperands())) {
       if (!definingOp || !definingOp.hasTensorSemantics())
         continue;
       auto fusedOp =
           fuseTensorOps(rewriter, cast<LinalgOp>(definingOp.getOperation()),
                         cast<LinalgOp>(op.getOperation()), operand.index());
       if (!fusedOp)
         continue;
       rewriter.replaceOp(op, fusedOp.getValue().getOperation()->getResults());
+      if (llvm::all_of(definingOp.getResults(),
+                       [](Value val) -> bool { return val.use_empty(); }))
+        rewriter.eraseOp(definingOp);
       return success();
     }
     return failure();
   }
 };
 
 /// Pass that fuses generic ops on tensors. Used only for testing.
 struct FusionOfTensorOpsPass

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

 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/LinalgTransforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/LoopOps/EDSC/Builders.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/Support/Functional.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/STLExtras.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
LinalgOpToLoopsImpl<LoopTy, ConcreteOpTy>::doit(Operation *op,
   assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   auto nPar = linalgOp.getNumParallelLoops();
   auto nRed = linalgOp.getNumReductionLoops();
   auto nWin = linalgOp.getNumWindowLoops();
   auto nLoops = nPar + nRed + nWin;
   auto mapsRange =
       linalgOp.indexing_maps().template getAsRange<AffineMapAttr>();
-  auto maps = llvm::to_vector<8>(
-      llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
+  auto maps =
+      functional::map([](AffineMapAttr a) { return a.getValue(); }, mapsRange);
   AffineMap invertedMap = inversePermutation(concatAffineMaps(maps));
+  if (!invertedMap)
+    return {};
   if (invertedMap.isEmpty()) {
     LinalgScopedEmitter<IndexedValueTy, ConcreteOpTy>::emitScalarImplementation(
         {}, linalgOp);
     return LinalgLoops();
   }
 
   SmallVector<ValueHandle, 4> allIvs(nLoops,
                                      ValueHandle(rewriter.getIndexType()));

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

LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: Permute dims for linalg op: " << *op
                     << ":\n");
 
   assert(succeeded(permuteGenericLinalgOpPrecondition(op, permutation)) &&
          "DRR failure case must be a precondition");
 
   auto linOp = cast<LinalgOp>(op);
   auto permutationMap = inversePermutation(
       AffineMap::getPermutationMap(permutation, rewriter.getContext()));
+  assert(permutationMap && "expected permutation to be invertible");
   SmallVector<AffineMap, 4> newIndexingMap;
   auto indexingMaps = linOp.indexing_maps().getValue();
   for (unsigned i = 0, e = linOp.getNumInputsAndOutputs(); i != e; ++i) {
-    AffineMap m = indexingMaps[i].cast<AffineMapAttr>().getValue().compose(
-        permutationMap);
+    AffineMap m = indexingMaps[i].cast<AffineMapAttr>().getValue();
+    if (!permutationMap.isEmpty())
+      m = m.compose(permutationMap);
     newIndexingMap.push_back(m);
   }
   auto itTypes = linOp.iterator_types().getValue();
   SmallVector<Attribute, 4> itTypesVector;
   for (unsigned i = 0, e = itTypes.size(); i != e; ++i)
     itTypesVector.push_back(itTypes[i]);
   applyPermutationToVector(itTypesVector, permutation);
   op->setAttr(getIndexingMapsAttrName(),

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

Optional<TiledLinalgOp> static tileLinalgOpImpl(OpBuilder &b, LinalgOp op,
 
   // If permutation is empty, use the identity. Build the permutation map
   // otherwise.
   auto invPermutationMap = AffineMap::getMultiDimIdentityMap(
       tileSizes.size(), ScopedContext::getContext());
   if (!permutation.empty())
     invPermutationMap = inversePermutation(
         AffineMap::getPermutationMap(permutation, ScopedContext::getContext()));
+  if (!invPermutationMap)
+    return llvm::None;
 
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(op);
   ScopedContext scope(b, op.getLoc());
   // 2. Build the tiled loop ranges.
   auto viewSizes = getViewSizes(b, op);
   // The flattened loopToOperandRangesMaps is expected to be an invertible
   // permutation map (asserted in the inverse calculation).
   auto mapsRange = op.indexing_maps().getAsRange<AffineMapAttr>();
   auto maps = llvm::to_vector<8>(
       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
   auto viewSizesToLoopsMap = inversePermutation(concatAffineMaps(maps));
-  assert(viewSizesToLoopsMap && "expected invertible map");
+  if (!viewSizesToLoopsMap)
+    return llvm::None;
 
   SmallVector<SubViewOp::Range, 4> loopRanges;
   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
   std::tie(loopRanges, loopIndexToRangeIndex) =
       makeTiledLoopRanges(b, scope.getLocation(), viewSizesToLoopsMap,
                           viewSizes, tileSizes, folder);
   if (!permutation.empty())
     applyPermutationToVector(loopRanges, permutation);

mlir/test/Dialect/Linalg/fusion-tensor.mlir

%2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %arg2 {
     }: tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
   return %2 : tensor<?x?xf32>
 }
 
 // -----
 
 // CHECK-DAG: [[MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-DAG: [[MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0)>
-// CHECK-DAG: [[MAP2:#[a-zA-Z0-9_]*]] = affine_map<(d0) -> (d0)>
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> (d0)>
 #map2 = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: @add_broadcast_mul_fusion
 func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
 {
   %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]} %arg0, %arg1 {

 // CHECK-LABEL: @add_mul_scalar_fusion
 func @add_mul_scalar_fusion(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32>
 {
   %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = []} %arg0, %arg1 {
     ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
       %1 = addf %arg3, %arg4 : f32
       linalg.yield %1 : f32
   }: tensor<f32>, tensor<f32> -> tensor<f32>
-  // CHECK: linalg.generic {args_in = 2 : i64, args_out = 1 : i64
+  // CHECK: linalg.generic {args_in = 3 : i64, args_out = 1 : i64
   // CHECK: addf
   // CHECK: mulf
   %1 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = []} %0, %arg2 {
     ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
       %1 = mulf %arg3, %arg4 : f32
       linalg.yield %1 : f32
   }: tensor<f32>, tensor<f32> -> tensor<f32>
 
   return %1 : tensor<f32>
 }