[mlir][Linalg] NFC - Refactor vectorization to be more composable

[bkramer]

looks good

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

 struct LinalgFusionOptions;
 struct LinalgTilingOptions;
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
 using LinalgLoops = SmallVector<Operation *, 4>;
 
-struct TiledLinalgOp {
-  LinalgOp op;
-  SmallVector<Operation *, 8> loops;
-  SmallVector<Value, 4> tensorResults;
-  TiledLinalgOp &operator=(const TiledLinalgOp &) = default;
-};
-
 /// Populates patterns for vectorization of all ConvN-D ops.
 void populateConvVectorizationPatterns(
     MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns,
     ArrayRef<int64_t> tileSizes);
 
 /// Populates the given list with patterns to bufferize linalg ops.
 void populateLinalgBufferizePatterns(MLIRContext *context,
                                      BufferizeTypeConverter &converter,

 ///
 /// Returns a struct containing the tiled loops in the specified order
 /// and the cloned op if successful, llvm::None otherwise.
 ///
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed by
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`tileSizes.size()` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
+struct TiledLinalgOp {
+  LinalgOp op;
+  SmallVector<Operation *, 8> loops;
+  SmallVector<Value, 4> tensorResults;
+  TiledLinalgOp &operator=(const TiledLinalgOp &) = default;
+};
 Optional<TiledLinalgOp> tileLinalgOp(OpBuilder &b, LinalgOp op,
                                      const LinalgTilingOptions &options);
 
 /// Fuse a sequence of linalg operations (`ops`) using tile-and-fuse. This
 /// proceeds as follows:
 /// - Find outer parallel loops in these ops that can be fused.
 /// - Tile fusable outer parallel loops of the last operation in the sequence.
 /// - Fuse the remaining operations with the tiled operation

 ///
 /// 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,
                                    OperationFolder *folder = nullptr);
 
 /// Emit a suitable vector form for a Linalg op with fully static shape.
-void vectorizeLinalgOp(OpBuilder &builder, Operation *op);
+struct VectorizedLinalgOp {
+  SmallVector<Value> tensorResults;
+  VectorizedLinalgOp &operator=(const VectorizedLinalgOp &) = default;
+};
+Optional<VectorizedLinalgOp> vectorizeLinalgOp(OpBuilder &builder,
+                                               Operation *op);
 
 /// Emits a loop nest of `LoopTy` with the proper body for `op`.
 template <typename LoopTy>
 Optional<LinalgLoops>
 linalgLowerOpToLoops(OpBuilder &builder, Operation *op,
                      ArrayRef<unsigned> interchangeVector = {});
 
 /// Emits a loop nest of `scf.for` with the proper body for `op`. The generated

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

 
 LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
-  if (failed(vectorizeLinalgOpPrecondition(op)))
+  Optional<VectorizedLinalgOp> res = vectorizeLinalgOp(rewriter, op);
+  if (!res)
     return failure();
-  vectorizeLinalgOp(rewriter, op);
+  if (!res->tensorResults.empty())
+    rewriter.replaceOp(op, res->tensorResults);
+  else
     rewriter.eraseOp(op);
   return success();
 }
 
 LogicalResult mlir::linalg::applyStagedPatterns(
     Operation *op, ArrayRef<FrozenRewritePatternList> stage1Patterns,
     const FrozenRewritePatternList &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda) {
   unsigned iteration = 0;

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

 ///   1. The region for the linalg op is created if necessary.
 ///   2. Values defined above the region are mapped to themselves and will be
 ///   broadcasted on a per-need basis by their consumers.
 ///   3. Each region argument is vectorized into a vector.transfer_read (or 0-d
 ///   load).
 ///   TODO: Reuse opportunities for RAR dependencies.
 ///   4. Register CustomVectorizationHook for YieldOp to capture the results.
 ///   5. Iteratively call vectorizeOneOp on the region operations.
-///   6. RAUW the linalg op by the results captured vectorizing the YieldOp.
-static LogicalResult vectorizeAsLinalgGeneric(
+static Optional<VectorizedLinalgOp> vectorizeAsLinalgGeneric(
     OpBuilder &builder, LinalgOp linalgOp,
     ArrayRef<CustomVectorizationHook> customVectorizationHooks = {}) {
   // 1. Certain Linalg ops do not have a region but only a region builder.
   // If so, build the region so we can vectorize.
   std::unique_ptr<Region> owningRegion;
   Region *region;
   if (linalgOp->getNumRegions() > 0) {
     region = &linalgOp->getRegion(0);
static Optional<VectorizedLinalgOp> vectorizeAsLinalgGeneric(
   auto hooks = llvm::to_vector<4>(customVectorizationHooks);
   hooks.push_back(vectorizeYield);
 
   // 5. Iteratively call `vectorizeOneOp` to each op in the slice.
   for (Operation &op : block->getOperations()) {
     VectorizationResult result = vectorizeOneOp(builder, &op, bvm, hooks);
     if (result.status == VectorizationStatus::Failure) {
       LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: failed to vectorize: " << op);
-      return failure();
+      return llvm::None;
     }
     if (result.status == VectorizationStatus::NewOp) {
       LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: new vector op: "
                         << *result.newOp;);
       bvm.map(op.getResults(), result.newOp->getResults());
     }
   }
 
-  // 6. RAUW the linalg op by the results captured vectorizing the YieldOp.
-  if (!results.empty())
-    linalgOp->replaceAllUsesWith(results);
-  return success();
+  return VectorizedLinalgOp{{results}};
 }
 
 /// Detect whether `r` has only ConstantOp, ElementwiseMappable and YieldOp.
 static bool hasOnlyScalarElementwiseOp(Region &r) {
   if (!llvm::hasSingleElement(r))
     return false;
   for (Operation &op : r.front()) {
     if (!(isa<ConstantOp, linalg::YieldOp>(op) ||
static bool isElementwise(Operation *op) {
   // to the right shape.
   for (unsigned i = 0, e = genericOp.getNumInputs(); i < e; i++) {
     if (!genericOp.getInputIndexingMap(i).isMinorIdentity())
       return false;
   }
   return hasOnlyScalarElementwiseOp(genericOp.getRegion());
 }
 
-static void vectorizeContraction(OpBuilder &builder, LinalgOp linalgOp) {
+static Optional<VectorizedLinalgOp> vectorizeContraction(OpBuilder &builder,
+                                                         LinalgOp linalgOp) {
   assert(isaContractionOpInterface(linalgOp) &&
          "expected vectorizeContraction preconditions to be met");
   Location loc = linalgOp.getLoc();
   // Vectorize other ops as vector contraction.
   // TODO: interface.
   LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: "
                     << "Rewrite linalg op as vector.contract: ";
              linalgOp.dump());
auto vType = outShape.empty()
                      : VectorType::get(outShape, op->getResult(0).getType());
     auto zero =
         builder.create<ConstantOp>(loc, vType, builder.getZeroAttr(vType));
     Operation *contract = builder.create<vector::ContractionOp>(
         loc, bvm.lookup(op->getOperand(0)), bvm.lookup(op->getOperand(1)), zero,
         linalgOp.indexing_maps(), linalgOp.iterator_types());
     return VectorizationResult{VectorizationStatus::NewOp, contract};
   };
-  auto status =
-      vectorizeAsLinalgGeneric(builder, linalgOp, {vectorizeContraction});
-  (void)status;
-  assert(succeeded(status) &&
-         "Unexpected vectorization failed despite preconditions");
+  return vectorizeAsLinalgGeneric(builder, linalgOp, {vectorizeContraction});
 }
 
 LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
   auto linalgOp = cast<linalg::LinalgOp>(op);
   // All types must be static shape to go to vector.
   for (Value operand : linalgOp.getShapedOperands())
     if (!operand.getType().cast<ShapedType>().hasStaticShape())
       return failure();
   for (Type outputTensorType : linalgOp.getOutputTensorTypes())
     if (!outputTensorType.cast<ShapedType>().hasStaticShape())
       return failure();
 
   if (isa<linalg::FillOp, linalg::CopyOp>(op))
     return success();
   if (isElementwise(op))
     return success();
   return success(isaContractionOpInterface(linalgOp));
 }
 
-void mlir::linalg::vectorizeLinalgOp(OpBuilder &builder, Operation *op) {
-  assert(succeeded(vectorizeLinalgOpPrecondition(op)));
+Optional<VectorizedLinalgOp> mlir::linalg::vectorizeLinalgOp(OpBuilder &builder,
+                                                             Operation *op) {
+  if (failed(vectorizeLinalgOpPrecondition(op)))
+    return llvm::None;
 
   edsc::ScopedContext scope(builder, op->getLoc());
   // In the case of 0-D memrefs, return null and special case to scalar load or
   // store later.
   if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
     // Vectorize fill as a vector.broadcast.
     LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: "
                       << "Rewrite linalg.fill as vector.broadcast: " << *op);
-    buildVectorWrite(builder, fillOp.value(), fillOp.output());
-    return;
+    VectorizedLinalgOp res;
+    if (Value v = buildVectorWrite(builder, fillOp.value(), fillOp.output()))
+      res.tensorResults.push_back(v);
+    return res;
   }
   if (auto copyOp = dyn_cast<linalg::CopyOp>(op)) {
     // Vectorize copy as a vector.transfer_read+vector.transfer_write.
     LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: "
                       << "Rewrite linalg.copy as vector.transfer_read + "
                          "vector.transfer_write: "
                       << *op);
     Value vector = buildVectorRead(builder, copyOp.input());
-    buildVectorWrite(builder, vector, copyOp.output());
-    return;
+    VectorizedLinalgOp res;
+    if (Value v = buildVectorWrite(builder, vector, copyOp.output()))
+      res.tensorResults.push_back(v);
+    return res;
   }
-
   if (isElementwise(op)) {
     LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: "
-                      << "Rewrite linalg op as vector.transfer_read + " << *op);
-    auto status = vectorizeAsLinalgGeneric(builder, cast<LinalgOp>(op));
-    (void)status;
-    assert(succeeded(status) &&
-           "Unexpected vectorization failed despite preconditions");
-    return;
+                      << "Vectorize linalg op as a generic: " << *op);
+    return vectorizeAsLinalgGeneric(builder, cast<LinalgOp>(op));
   }
 
-  vectorizeContraction(builder, cast<LinalgOp>(op));
+  // TODO: as soon as Copy and FillOp. get a region builder, replace all the
+  // above by:
+  // if (isa<FillOp, CopyOp>(op) || isElementwise(op)) {
+  //   LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: "
+  //                     << "Vectorize linalg op as a generic: " << *op);
+  //   return vectorizeAsLinalgGeneric(builder, cast<LinalgOp>(op));
+  // }
+
+  return vectorizeContraction(builder, cast<LinalgOp>(op));
 }
 
 //----------------------------------------------------------------------------//
 // Misc. conv vectorization patterns.
 //----------------------------------------------------------------------------//
 // TODO: cleanup all this.
 template <class ConvOp, int N>
 LogicalResult ConvOpVectorization<ConvOp, N>::matchAndRewrite(

mlir/test/Dialect/Linalg/vectorization.mlir

-// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns -split-input-file -debug-only=linalg-vectorization
+
+//| FileCheck %s
 
 // -----
 
 // CHECK-LABEL: contraction_dot
 func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
   // CHECK: vector.contract
   // CHECK-SAME: vector<1584xf32>, vector<1584xf32> into f32
   linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>)