[mlir][Linalg] NFC - Drop Optional in favor of FailureOr

Differential revision: https://reviews.llvm.org/D112332

Author: nicolasvasilache

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

@@ -158,8 +158,8 @@ struct TiledLinalgOp {
   SmallVector<Operation *, 8> loops;
   SmallVector<Value, 4> tensorResults;
 };
-Optional<TiledLinalgOp> tileLinalgOp(OpBuilder &b, LinalgOp op,
-                                     const LinalgTilingOptions &options);
+FailureOr<TiledLinalgOp> tileLinalgOp(OpBuilder &b, LinalgOp op,
+                                      const LinalgTilingOptions &options);
 
 /// Fuse a sequence of linalg operations (`ops`) using tile-and-fuse. This
 /// proceeds as follows:
@@ -221,7 +221,7 @@ struct TiledAndFusedLinalgOps {
   /// The fused loop generated.
   SmallVector<Operation *, 4> fusedLoops;
 };
-Optional<TiledAndFusedLinalgOps>
+FailureOr<TiledAndFusedLinalgOps>
 tileAndFuseLinalgOps(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                      const LinalgDependenceGraph &dependenceGraph,
                      const LinalgTilingOptions &tilingOptions);
@@ -344,7 +344,7 @@ struct PromotionInfo {
   Value fullLocalView;
   Value partialLocalView;
 };
-Optional<PromotionInfo>
+FailureOr<PromotionInfo>
 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
                           AllocBufferCallbackFn allocationFn,
                           DataLayout &layout);
@@ -359,24 +359,24 @@ promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
 ///
 /// Returns the modified linalg op (the modification happens in place) as well
 /// as all the copy ops created.
-Optional<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
-                                   LinalgPromotionOptions options);
+FailureOr<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
+                                    LinalgPromotionOptions options);
 
 /// Emit a suitable vector form for a Linalg op with fully static shape.
 LogicalResult vectorizeLinalgOp(OpBuilder &builder, Operation *op,
                                 SmallVectorImpl<Value> &newResults);
 
 /// Emits a loop nest of `scf.for` with the proper body for `linalgOp`.
-Optional<LinalgLoops> linalgOpToLoops(PatternRewriter &rewriter,
-                                      LinalgOp linalgOp);
+FailureOr<LinalgLoops> linalgOpToLoops(PatternRewriter &rewriter,
+                                       LinalgOp linalgOp);
 
 /// Emits a loop nest of `scf.parallel` with the proper body for `linalgOp`.
-Optional<LinalgLoops> linalgOpToParallelLoops(PatternRewriter &rewriter,
-                                              LinalgOp linalgOp);
+FailureOr<LinalgLoops> linalgOpToParallelLoops(PatternRewriter &rewriter,
+                                               LinalgOp linalgOp);
 
 /// Emits a loop nest of `affine.for` with the proper body for `linalgOp`.
-Optional<LinalgLoops> linalgOpToAffineLoops(PatternRewriter &rewriter,
-                                            LinalgOp linalgOp);
+FailureOr<LinalgLoops> linalgOpToAffineLoops(PatternRewriter &rewriter,
+                                             LinalgOp linalgOp);
 
 //===----------------------------------------------------------------------===//
 // Preconditions that ensure the corresponding transformation succeeds and can
@@ -961,15 +961,15 @@ struct LinalgLoweringPattern : public RewritePattern {
       // TODO: Move lowering to library calls here.
       return failure();
     case LinalgLoweringType::Loops:
-      if (!linalgOpToLoops(rewriter, op))
+      if (failed(linalgOpToLoops(rewriter, op)))
         return failure();
       break;
     case LinalgLoweringType::AffineLoops:
-      if (!linalgOpToAffineLoops(rewriter, op))
+      if (failed(linalgOpToAffineLoops(rewriter, op)))
         return failure();
       break;
     case LinalgLoweringType::ParallelLoops:
-      if (!linalgOpToParallelLoops(rewriter, op))
+      if (failed(linalgOpToParallelLoops(rewriter, op)))
         return failure();
       break;
     }

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

@@ -164,25 +164,25 @@ struct FusionInfo {
 /// Implements the fusion part of the "tileAndFuse on buffers" transformation
 /// and thus requires the `consumerOpOperand` to be a `subview` op (generally
 /// obtained by applying the tiling transformation).
-Optional<FusionInfo> fuseProducerOfBuffer(OpBuilder &b,
-                                          OpOperand &consumerOpOperand,
-                                          const LinalgDependenceGraph &graph);
+FailureOr<FusionInfo> fuseProducerOfBuffer(OpBuilder &b,
+                                           OpOperand &consumerOpOperand,
+                                           const LinalgDependenceGraph &graph);
 /// Tensor counterpart of `fuseProducerOfBuffer`.
 /// This implements the fusion part of the "tileAndFuse on tensors"
 /// transformation and thus requires the `consumerOpOperand` to be a
 /// `extract_slice` op (generally obtained by applying the tiling
 /// transformation).
-Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
-                                          OpOperand &consumerOpOperand);
+FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
+                                           OpOperand &consumerOpOperand);
 /// Tensor counterpart of `fuseProducerOfBuffer`.
 /// This implements the fusion part of the "tileAndFuse on tensors"
 /// transformation and thus requires the `consumerOpOperand` to be a
 /// `extract_slice` op (generally obtained by applying the tiling
 /// transformation). Assumes `producerOfTensor` is a Linalg op that produces
 /// `consumerOpOperand`.
-Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
-                                          OpResult producerOpResult,
-                                          OpOperand &consumerOpOperand);
+FailureOr<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
+                                           OpResult producerOpResult,
+                                           OpOperand &consumerOpOperand);
 
 //===----------------------------------------------------------------------===//
 // Fusion on tensor utilities

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

@@ -331,7 +331,7 @@ bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph,
 /// For `consumer` with buffer semantics, find the Linalg operation on buffers
 /// that is the last writer of `consumerOpOperand`. For now the fusable
 /// dependence is returned as an instance of the `dependenceGraph`.
-static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
+static FailureOr<LinalgDependenceGraph::LinalgDependenceGraphElem>
 findFusableProducer(OpOperand &consumerOpOperand,
                     const LinalgDependenceGraph &dependenceGraph) {
   LLVM_DEBUG(llvm::dbgs() << "findFusableProducer for: "
@@ -340,7 +340,7 @@ findFusableProducer(OpOperand &consumerOpOperand,
                           << *consumerOpOperand.getOwner() << "\n");
   LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner());
   if (!consumerOp)
-    return {};
+    return failure();
 
   // Only consider RAW and WAW atm.
   for (auto depType : {
@@ -386,37 +386,37 @@ findFusableProducer(OpOperand &consumerOpOperand,
       }
     }
   }
-  return {};
+  return failure();
 }
 
-Optional<FusionInfo>
+FailureOr<FusionInfo>
 mlir::linalg::fuseProducerOfBuffer(OpBuilder &b, OpOperand &consumerOpOperand,
                                    const LinalgDependenceGraph &graph) {
   Optional<LinalgDependenceGraph::LinalgDependenceGraphElem> fusableDependence =
       findFusableProducer(consumerOpOperand, graph);
   if (!fusableDependence)
-    return llvm::None;
+    return failure();
 
   LinalgOp producerOp = dyn_cast<LinalgOp>(fusableDependence->getDependentOp());
   if (!producerOp)
-    return llvm::None;
+    return failure();
 
   // If producer is already in the same block as consumer, we are done.
   if (consumerOpOperand.get().getParentBlock() ==
       fusableDependence->getDependentValue().getParentBlock())
-    return llvm::None;
+    return failure();
 
   Optional<AffineMap> producerMap =
       fusableDependence->getDependentOpViewIndexingMap();
   if (!producerMap)
-    return llvm::None;
+    return failure();
 
   // Must be a subview or an extract_slice to guarantee there are loops we can
   // fuse into.
   auto subView = consumerOpOperand.get().getDefiningOp<memref::SubViewOp>();
   if (!subView) {
     LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subview)");
-    return llvm::None;
+    return failure();
   }
 
   // Fuse `producer` just before `consumer`.
@@ -459,28 +459,28 @@ static void getProducerOfTensor(Value tensor, OpResult &opResult) {
   }
 }
 
-Optional<FusionInfo>
+FailureOr<FusionInfo>
 mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
   Value inputTensor = consumerOpOperand.get();
   OpResult producerOpResult;
   getProducerOfTensor(inputTensor, producerOpResult);
   if (!producerOpResult) {
     LLVM_DEBUG(llvm::dbgs() << "\nUnable to find producer");
-    return {};
+    return failure();
   }
   return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
 }
 
-Optional<FusionInfo>
+FailureOr<FusionInfo>
 mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
                                    OpOperand &consumerOpOperand) {
   auto producerOp = dyn_cast<LinalgOp>(producerOpResult.getOwner());
   if (!producerOp)
-    return llvm::None;
+    return failure();
 
   LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner());
   if (!consumerOp)
-    return llvm::None;
+    return failure();
 
   Value inputTensor = consumerOpOperand.get();
 
@@ -489,13 +489,13 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
   if (!sliceOp) {
     LLVM_DEBUG(llvm::dbgs()
                << "\nNot fusable, not an extract_slice op: " << inputTensor);
-    return {};
+    return failure();
   }
 
   // If producer is already in the same block as consumer, we are done.
   if (consumerOpOperand.get().getParentBlock() ==
       producerOpResult.getParentBlock())
-    return {};
+    return failure();
 
   // Insert fused `producer` just before `consumer`.
   OpBuilder::InsertionGuard g(b);
@@ -537,27 +537,27 @@ static AffineMap pruneReductionDimsFromMap(ArrayRef<Attribute> iteratorTypes,
 /// - indexing map of the fused view in the producer : producerIndexMap
 ///     consumerLoopToProducerLoop =
 ///       inverse(producerIndexMap).compose(consumerIndexMap)
-static Optional<AffineMap> getConsumerLoopToProducerLoopMap(
+static FailureOr<AffineMap> getConsumerLoopToProducerLoopMap(
     LinalgDependenceGraph::LinalgDependenceGraphElem dependence) {
   auto producer = dyn_cast<LinalgOp>(dependence.getDependentOp());
   if (!producer)
-    return None;
+    return failure();
 
   Optional<AffineMap> producerIndexingMap =
       dependence.getDependentOpViewIndexingMap();
   Optional<AffineMap> consumerIndexingMap =
       dependence.getIndexingOpViewIndexingMap();
   if (!producerIndexingMap || !consumerIndexingMap)
-    return None;
+    return failure();
 
   AffineMap prunedProducerIndexingMap = pruneReductionDimsFromMap(
       producer.iterator_types().getValue(), *producerIndexingMap);
   if (!prunedProducerIndexingMap.isPermutation())
-    return None;
+    return failure();
 
   if (consumerIndexingMap->getNumResults() !=
       prunedProducerIndexingMap.getNumResults())
-    return None;
+    return failure();
 
   LLVM_DEBUG({
     llvm::dbgs() << "\t producerMap : ";
@@ -572,7 +572,7 @@ static Optional<AffineMap> getConsumerLoopToProducerLoopMap(
 
   AffineMap invProducerIndexMap = inversePermutation(prunedProducerIndexingMap);
   if (!invProducerIndexMap)
-    return None;
+    return failure();
 
   return invProducerIndexMap.compose(*consumerIndexingMap);
 }
@@ -776,7 +776,7 @@ FusableOpDependencesTy mlir::linalg::findAllFusableDependences(
 
 /// Tile the fused loops in the root operation, by setting the tile sizes for
 /// all other loops to zero (those will be tiled later).
-static Optional<TiledLinalgOp>
+static FailureOr<TiledLinalgOp>
 tileRootOperation(OpBuilder &b, LinalgOp op, ArrayRef<Value> tileSizeVector,
                   const LinalgTilingOptions &options,
                   const std::set<unsigned> &fusedLoops) {
@@ -871,12 +871,12 @@ fuseOperations(OpBuilder &b, LinalgOp rootOp, TiledLinalgOp tiledLinalgOp,
   return fusedOps;
 }
 
-static Optional<TiledAndFusedLinalgOps>
+static FailureOr<TiledAndFusedLinalgOps>
 tileAndFuseLinalgOpsImpl(OpBuilder &b, ArrayRef<LinalgOp> ops,
                          const LinalgDependenceGraph &dependenceGraph,
                          const LinalgTilingOptions &tilingOptions) {
   if (ops.size() < 2)
-    return llvm::None;
+    return failure();
   LinalgOp rootOp = ops.back();
   if (!llvm::all_of(
           ops,
@@ -887,13 +887,13 @@ tileAndFuseLinalgOpsImpl(OpBuilder &b, ArrayRef<LinalgOp> ops,
     rootOp.emitError(
         "unable to fuse operations that have tensor semantics with operations "
         "that have buffer semantics and viceversa.");
-    return llvm::None;
+    return failure();
   }
   // TODO: Support interchange with tile + fuse. This might actually help do
   // better fusion.
   if (!tilingOptions.interchangeVector.empty()) {
     rootOp.emitRemark("unable to handle tile and fuse with interchange");
-    return llvm::None;
+    return failure();
   }
 
   OpBuilder::InsertionGuard guard(b);
@@ -905,7 +905,7 @@ tileAndFuseLinalgOpsImpl(OpBuilder &b, ArrayRef<LinalgOp> ops,
       findAllFusableDependences(ops, dependenceGraph);
   if (fusableDependences.empty()) {
     LLVM_DEBUG(llvm::dbgs() << "no fusable dependencies found\n");
-    return llvm::None;
+    return failure();
   }
 
   TiledAndFusedLinalgOps ret;
@@ -917,17 +917,17 @@ tileAndFuseLinalgOpsImpl(OpBuilder &b, ArrayRef<LinalgOp> ops,
   // just return.
   if (ret.fusedLoopDims.empty()) {
     LLVM_DEBUG(llvm::dbgs() << "no fusable loops found\n");
-    return llvm::None;
+    return failure();
   }
 
   // Tile the fused loops in the last operation in the list.
   SmallVector<Value, 4> tileSizeVector =
       tilingOptions.tileSizeComputationFunction(b, rootOp);
-  Optional<TiledLinalgOp> tiledRootOp = tileRootOperation(
+  FailureOr<TiledLinalgOp> tiledRootOp = tileRootOperation(
       b, rootOp, tileSizeVector, tilingOptions, ret.fusedLoopDims);
-  if (!tiledRootOp) {
+  if (failed(tiledRootOp)) {
     rootOp.emitRemark("failed to tile the fused loops");
-    return llvm::None;
+    return failure();
   }
   ret.op = tiledRootOp->op;
   ret.fusedLoops.assign(tiledRootOp->loops.begin(), tiledRootOp->loops.end());
@@ -939,7 +939,7 @@ tileAndFuseLinalgOpsImpl(OpBuilder &b, ArrayRef<LinalgOp> ops,
   return ret;
 }
 
-Optional<TiledAndFusedLinalgOps>
+FailureOr<TiledAndFusedLinalgOps>
 mlir::linalg::tileAndFuseLinalgOps(OpBuilder &b, ArrayRef<LinalgOp> ops,
                                    const LinalgDependenceGraph &dependenceGraph,
                                    const LinalgTilingOptions &tilingOptions) {
@@ -950,5 +950,5 @@ mlir::linalg::tileAndFuseLinalgOps(OpBuilder &b, ArrayRef<LinalgOp> ops,
     return tileAndFuseLinalgOpsImpl(b, ops, dependenceGraph, tilingOptions);
   default:;
   }
-  return llvm::None;
+  return failure();
 }