[mlir][linalg] Verify every LinalgOp has a body.

[gysit]

Summary

After removing the last LinalgOps that have no region attached we can verify there is a region. The patch performs the following changes:

• Move the SingleBlockImplicitTerminator trait further up the the structured op base class.
• Adapt the LinalgOp verification since the trait only check if there is 0 or 1 block.
• Introduce a getBlock method on the LinalgOp interface.
• Access the LinalgOp body using either getBlock() or getBody() if the concrete operation type is known.

This patch is a follow up to https://reviews.llvm.org/D111233.

[gysit]

Rebase.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
         uninitialized, as well as properly cloning "read-write" operands.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"payloadUsesValueFromOperand",
       /*args=*/(ins "OpOperand *":$opOperand),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         unsigned bbArgNumber = opOperand->getOperandNumber();
-        // Safeguard against the named linalg ops that are manually defined and
-        // that only support buffer semantics: we should not be there.
-        // Such ops have an empty regionBuilder and are not constructed with a
-        // region for now. In the future they are slated to disappear.
-        assert(this->getOperation()->getNumRegions() == 1 && "unexpected "
-               "missing region (calling `payloadUsesValueFromOperand` on "
-               "manually defined named Linalg op?)");
-        Block &block = this->getOperation()->getRegion(0).front();
         // Init tensors have uses.
-        return !block.getArgument(bbArgNumber).use_empty();
+        return !getBlock()->getArgument(bbArgNumber).use_empty();
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return true if `opOperand` is an input tensor.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"isInputTensor",
InterfaceMethod<
       /*desc=*/[{
         Return the output block arguments of the region.
       }],
       /*retTy=*/"Block::BlockArgListType",
       /*methodName=*/"getRegionOutputArgs",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        Block &entryBlock = this->getOperation()->getRegion(0).front();
-        return entryBlock.getArguments().take_back(this->getNumOutputs());
+        return getBlock()->getArguments().take_back(this->getNumOutputs());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the `opOperand` shape or an empty vector for scalars.
       }],
       /*retTy=*/"ArrayRef<int64_t>",
       /*methodName=*/"getShape",
InterfaceMethod<
         return this->getOperation()->getResult(resultIndex);
       }]
     >,
     //===------------------------------------------------------------------===//
     // Other interface methods.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/[{
+        Return the single block constituting the body of the operation by
+        calling the getBody method on the concrete operation.
+      }],
+      /*retTy=*/"Block*",
+      /*methodName=*/"getBlock",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        // Assume the concrete operation implements the
+        // SingleBlockImplicitTerminator trait.
+        return $_op.getBody();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the iterator types attribute within the current operation.
       }],
       /*retTy=*/"ArrayAttr",
       /*methodName=*/"iterator_types",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.iterator_types();

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // Base Tablegen class for Linalg ops.
 // Linalg ops that correspond to library calls operate on ShapedType as their
 // first operands. These may be optionally followed by non-view operands
 // depending on the specific Linalg op.
 class LinalgStructuredBase_Op<string mnemonic, list<OpTrait> props>
   : Op<Linalg_Dialect, mnemonic, !listconcat([
-       LinalgStructuredInterface, ReifyRankedShapedTypeOpInterface], props)> {
+       SingleBlockImplicitTerminator<"YieldOp">,
+       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       LinalgStructuredInterface,
+       ReifyRankedShapedTypeOpInterface], props)> {
   code structuredOpsBaseDecls = [{
     // Return whether the op accesses the iteration indices.
     bool hasIndexSemantics() {
-      Operation *op = this->getOperation();
-      if(op->getNumRegions() == 0 || op->getRegion(0).empty())
-        return false;
-      return !op->getRegion(0).front().getOps<IndexOp>().empty();
+      return !this->getBody()->getOps<IndexOp>().empty();
     }
 
     LogicalResult reifyResultShapes(OpBuilder &b,
         ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
       return cast<LinalgOp>(getOperation()).reifyResultShapes(b,
           reifiedReturnShapes);
     }
   }];
 }
 
 class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
-  : LinalgStructuredBase_Op<mnemonic,
-       !listconcat(props, [
-         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>])> {
+  : LinalgStructuredBase_Op<mnemonic, !listconcat(props, [])> {
   code structuredOpsDecls = structuredOpsBaseDecls # [{
     std::string getLibraryCallName() {
       return generateLibraryCallName(getOperation());
     }
   }];
   let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
 }
 
def FillOp : LinalgStructured_Op<"fill", []> {
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Generic Linalg ops.
 //===----------------------------------------------------------------------===//
 
-def GenericOp : LinalgStructuredBase_Op<"generic", [
-    AttrSizedOperandSegments,
-    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-    SingleBlockImplicitTerminator<"YieldOp">]> {
+def GenericOp : LinalgStructuredBase_Op<"generic", [AttrSizedOperandSegments]> {
   let description = [{
     Generic Linalg op form where the key properties of the computation are
     specified as attributes. In pretty form, a `linalg.generic` op is written
     as:
 
       ```mlir
       linalg.generic #trait_attribute
           ins(%A, %B : memref<?x?xf32, stride_specification>,

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

for (auto expr : indexingMap.getResults()) {
                  << (opOperand->getOperandNumber() - linalgOp.getNumInputs())
                  << " a.k.a '" << exprStr
                  << "' is function of reduction iterator 'd" << pos << "'";
         }
       }
     }
   }
 
-  // Named ops that are defined manually have a region builder but no region at
-  // this time. Assume the region is well-formed by specification.
-  // TODO: use linalg-ods-gen for all ops when we have enough expressive power.
-  if (linalgOp->getNumRegions() == 0) {
-    assert(!linalgOp.getRegionBuilder() && "regionBuilder but no region");
-    return success();
-  }
-
-  auto &region = linalgOp->getRegion(0);
-  if (linalgOp->getNumRegions() > 1 || !llvm::hasSingleElement(region))
-    return op->emitOpError("expected 1 region with 1 block");
+  // Check the region has exactly one block.
+  if (linalgOp->getNumRegions() != 1 ||
+      !llvm::hasSingleElement(linalgOp->getRegion(0)))
+    return op->emitOpError("expects to have 1 region with 1 block");
 
   if (!linalgOp.getShapesToLoopsMap())
     return op->emitOpError("expected the shape-to-loops map to be non-null");
 
   // Simplifying assumption: bbargs match 1-1 with shape operands elemental
   // types.
   // TODO: once ranked shape types are plugged in, we may want to drop the
   // corresponding bbargs, that can never be read from. This will be subject to

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

return ArrayAttr::get(context,
                               return AffineMapAttr::get(map);
                             })));
 }
 
 /// Update the index accesses of linalg operations having index semantics.
 static void replaceUnitDimIndexOps(GenericOp genericOp,
                                    const DenseSet<unsigned> &unitDims,
                                    PatternRewriter &rewriter) {
-  assert(genericOp->getNumRegions() == 1 &&
-         genericOp->getRegion(0).getBlocks().size() == 1 &&
-         "expected generic operation to have one block.");
-  Block &block = genericOp->getRegion(0).front();
-
-  for (IndexOp indexOp : llvm::make_early_inc_range(block.getOps<IndexOp>())) {
+  for (IndexOp indexOp :
+       llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
     OpBuilder::InsertionGuard guard(rewriter);
     rewriter.setInsertionPoint(indexOp);
     if (unitDims.count(indexOp.dim()) != 0) {
       rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(indexOp, 0);
     } else {
       // Update the dimension of the index operation if needed.
       unsigned droppedDims = llvm::count_if(
           unitDims, [&](unsigned dim) { return dim < indexOp.dim(); });

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

GenericOp mlir::linalg::generalizeNamedOp(PatternRewriter &rewriter,
   SmallVector<Value> inputOperands = namedOp.getInputOperands();
   SmallVector<Value> outputOperands = namedOp.getOutputOperands();
   SmallVector<AffineMap> indexingMaps = namedOp.getIndexingMaps();
   SmallVector<StringRef> iterators = llvm::to_vector<4>(
       namedOp.iterator_types().getAsValueRange<StringAttr>());
   SmallVector<RankedTensorType> resultTypes = namedOp.getOutputTensorTypes();
   SmallVector<Type> types(resultTypes.begin(), resultTypes.end());
 
-  // Inline the existing region if the named operation has a region attached.
-  if (namedOp->getNumRegions() == 1) {
-    GenericOp genericOp =
-        rewriter.create<GenericOp>(namedOp.getLoc(), types, inputOperands,
-                                   outputOperands, indexingMaps, iterators);
-    rewriter.inlineRegionBefore(namedOp->getRegion(0), genericOp.region(),
-                                genericOp.region().begin());
-    return genericOp;
-  }
+  // All named ops have a region attached that can be inlined.
+  assert(namedOp->getNumRegions() == 1 &&
+         "expect named op to have one region attached");
 
-  // Otherwise use the region builder to generate a new region.
-  // TODO: Remove this path once all linag operations have a region attached.
-  auto regionBuilder = namedOp.getRegionBuilder();
-  assert(regionBuilder && "expect the operation to have region builder");
-  return rewriter.create<GenericOp>(
-      namedOp.getLoc(), types, inputOperands, outputOperands, indexingMaps,
-      iterators,
-      [&regionBuilder](OpBuilder &bodyBuilder, Location loc, ValueRange) {
-        ImplicitLocOpBuilder b(loc, bodyBuilder);
-        regionBuilder(b, *bodyBuilder.getBlock());
-      });
-}
-
 namespace {
 
 struct LinalgGeneralizationPass
     : public LinalgGeneralizationBase<LinalgGeneralizationPass> {
   void runOnFunction() override;
 };
 
 } // namespace

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

void mlir::linalg::interchangeGenericOp(PatternRewriter &rewriter,
   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 &&
-           genericOp->getRegion(0).getBlocks().size() == 1 &&
-           "expected generic operation to have one block.");
-    Block &block = genericOp->getRegion(0).front();
     OpBuilder::InsertionGuard guard(rewriter);
     for (IndexOp indexOp :
-         llvm::make_early_inc_range(block.getOps<IndexOp>())) {
+         llvm::make_early_inc_range(genericOp.getBody()->getOps<IndexOp>())) {
       rewriter.setInsertionPoint(indexOp);
       SmallVector<Value> allIndices;
       allIndices.reserve(genericOp.getNumLoops());
       llvm::transform(llvm::seq<uint64_t>(0, genericOp.getNumLoops()),
                       std::back_inserter(allIndices), [&](uint64_t dim) {
                         return rewriter.create<IndexOp>(indexOp->getLoc(), dim);
                       });
       rewriter.replaceOpWithNewOp<AffineApplyOp>(
           indexOp, permutationMap.getSubMap(indexOp.dim()), allIndices);
     }
   }
 }

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

if (!linalgOp)
     return false;
   if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
     return false;
   // TODO: relax the restrictions on indexing map.
   for (OpOperand *opOperand : linalgOp.getOutputOperands()) {
     if (!linalgOp.getTiedIndexingMap(opOperand).isIdentity())
       return false;
   }
-  if (linalgOp->getNumRegions() != 1)
-    return false;
   return hasOnlyScalarElementwiseOp(linalgOp->getRegion(0));
 }
 
 /// Generic vectorization function that rewrites the body of a `linalgOp` into
 /// vector form. Generic vectorization proceeds as follows:
 ///   1. Verify the `linalgOp` has one non-empty region.
 ///   2. Values defined above the region are mapped to themselves and will be
 ///   broadcasted on a per-need basis by their consumers.

 /// reductions should occur, without any bookkeeping. The tradeoff is that, in
 /// the absence of good canonicalizations, the amount of work increases.
 /// This is not deemed a problem as we expect canonicalizations and foldings to
 /// aggressively clean up the useless work.
 LogicalResult vectorizeAsLinalgGeneric(
     OpBuilder &b, LinalgOp linalgOp, SmallVectorImpl<Value> &newResults,
     bool broadcastToMaximalCommonShape = false,
     ArrayRef<CustomVectorizationHook> customVectorizationHooks = {}) {
-  // 1. Fail to vectorize if the operation does not have one non-empty region.
-  if (linalgOp->getNumRegions() != 1 || linalgOp->getRegion(0).empty())
-    return failure();
-  auto &block = linalgOp->getRegion(0).front();
+  Block *block = linalgOp.getBlock();
 
   // 2. Values defined above the region can only be broadcast for now. Make them
   // map to themselves.
   BlockAndValueMapping bvm;
   SetVector<Value> valuesSet;
   mlir::getUsedValuesDefinedAbove(linalgOp->getRegion(0), valuesSet);
   bvm.map(valuesSet.getArrayRef(), valuesSet.getArrayRef());
 
   if (linalgOp.getNumOutputs() == 0)
     return failure();
 
   // TODO: the common vector shape is equal to the static loop sizes only when
   // all indexing maps are projected permutations. For convs and stencils the
   // logic will need to evolve.
   SmallVector<int64_t> commonVectorShape = linalgOp.computeStaticLoopSizes();
 
   // 3. Turn all BBArgs into vector.transfer_read / load.
   SmallVector<AffineMap> indexings;
   for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
-    BlockArgument bbarg = block.getArgument(opOperand->getOperandNumber());
+    BlockArgument bbarg = block->getArgument(opOperand->getOperandNumber());
     if (linalgOp.isScalar(opOperand)) {
       bvm.map(bbarg, opOperand->get());
       continue;
     }
     // TODO: 0-d vectors.
     Type readType;
     AffineMap map;
     if (linalgOp.getShape(opOperand).empty()) {
LogicalResult vectorizeAsLinalgGeneric(
   CustomVectorizationHook vectorizeIndex =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
     return vectorizeLinalgIndex(b, op, linalgOp);
   };
   hooks.push_back(vectorizeIndex);
 
   // 5. Iteratively call `vectorizeOneOp` to each op in the slice.
-  for (Operation &op : block.getOperations()) {
+  for (Operation &op : block->getOperations()) {
     VectorizationResult result = vectorizeOneOp(b, &op, bvm, hooks);
     if (result.status == VectorizationStatus::Failure) {
       LDBG("failed to vectorize: " << op);
       return failure();
     }
     if (result.status == VectorizationStatus::NewOp) {
       LDBG("new vector op: " << *result.newOp;);
       bvm.map(op.getResults(), result.newOp->getResults());

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

SmallVector<Value, 4> makeTiledShapes(OpBuilder &b, Location loc,
   }
 
   return tiledShapes;
 }
 
 void addTileLoopIvsToIndexOpResults(OpBuilder &b, LinalgOp tiledOp,
                                     ArrayRef<Value> ivs) {
   if (tiledOp.hasIndexSemantics()) {
-    assert(tiledOp->getNumRegions() == 1 &&
-           tiledOp->getRegion(0).getBlocks().size() == 1 &&
-           "expect producer to have one block.");
-    // Shift all IndexOp results by the tile offset.
-    Block &block = tiledOp->getRegion(0).front();
-    for (IndexOp indexOp : block.getOps<IndexOp>()) {
+    for (IndexOp indexOp : tiledOp.getBlock()->getOps<IndexOp>()) {
       if (ivs[indexOp.dim()] == nullptr)
         continue;
       OpBuilder::InsertionGuard guard(b);
       b.setInsertionPointAfter(indexOp);
       AffineExpr index, offset;
       bindDims(b.getContext(), index, offset);
       AffineApplyOp applyOp = makeComposedAffineApply(
           b, indexOp.getLoc(), index + offset,
           ValueRange{indexOp.getResult(), ivs[indexOp.dim()]});
       indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
     }
   }
 }
 
 } // namespace linalg
 } // namespace mlir

mlir/test/Dialect/Linalg/invalid.mlir

 ////////////////////////////////////////////////////////////////////////////////
 ///////////////////////////// Region tests /////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
 // -----
 
 func @generic_empty_region(%arg0: memref<f32>) {
   %f0 = arith.constant 0.0: f32
-  // expected-error @+1 {{op expected 1 region with 1 block}}
+  // expected-error @+1 {{op expects region #0 to have 0 or 1 blocks}}
   linalg.generic {
     indexing_maps =  [ affine_map<() -> ()>, affine_map<() -> ()> ],
     iterator_types = []}
       ins(%arg0 : memref<f32>)
      outs(%arg0 : memref<f32>) {
     ^bb1:
       linalg.yield %f0: f32
     ^bb2:
       linalg.yield %f0: f32
   }
 }
 
 // -----
 
 func @generic_empty_region(%arg0: memref<f32>) {
   %f0 = arith.constant 0.0: f32
-  // expected-error @+1 {{linalg.generic' op expected 1 region with 1 block}}
+  // expected-error @+1 {{op expects to have 1 region with 1 block}}
   linalg.generic {
     indexing_maps =  [ affine_map<() -> ()> , affine_map<() -> ()> ],
     iterator_types = []}
     ins(%arg0 : memref<f32>)
    outs(%arg0 : memref<f32>) {
   }
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

 def TestLinalgConvOpNotLinalgOp : TEST_Op<"conv_op_not_linalg_op", [
     LinalgConvolutionOpInterface]> {
   let arguments = (ins
     AnyType:$image, AnyType:$filter, AnyType:$output);
   let results = (outs AnyRankedTensor:$result);
 }
 
 def TestLinalgConvOp :
-  TEST_Op<"linalg_conv_op", [AttrSizedOperandSegments,
+  TEST_Op<"linalg_conv_op", [AttrSizedOperandSegments, SingleBlock,
       LinalgStructuredInterface, LinalgConvolutionOpInterface]> {
 
   let arguments = (ins Variadic<AnyType>:$inputs,
     Variadic<AnyType>:$outputs);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = [{

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

 // {4}: op attribute list
 // {5}: builder methods taking standalone attribute parameters
 // {6}: additional methods for attributes used by indexing maps
 static const char structuredOpOdsHeaderFormat[] = R"FMT(
 //===----------------------------------------------------------------------===//
 // Op definition for {0}
 //===----------------------------------------------------------------------===//
 
-def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([
-  AttrSizedOperandSegments,
-  DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-  SingleBlockImplicitTerminator<"YieldOp">],
+def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([AttrSizedOperandSegments],
   /*extraInterfaces=*/[{2}])> {
     {3}
     let arguments = (ins
       Variadic<AnyType>:$inputs,
       Variadic<AnyShaped>:$outputs{4}
     );
     let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
     let regions = (region AnyRegion:$region);