[mlir][sparse] integrate sparse annotation into generic linalg op

[aartbik]

Summary

This CL integrates the new sparse annotations (hereto merely added as fully
transparent attributes) more tightly to the generic linalg op in order to add
verification of the annotations' consistency as well as to make make other
passes more aware of their presence (in the long run, rewriting rules must
preserve the integrity of the annotations).

[nicolasvasilache]

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

496

Maybe add a comment that this is an ArrayAttr of StrArrayAttr?

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

441

Please factor out the magic string into mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h here and below

[jpienaar]

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

554

Are these the only two expected here? E.g., could we switch to enum rather than string attribute below?

[aartbik]

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

554

This is a bit forward looking in the sense that keeping this a string will simplify prototyping other storage formats that go beyond a simple enum (perhaps a property + parameters). So unless you foresee issues, I would prefer not to commit to a enum as of yet (other than what Nicolas proposed).

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

441

That makes sense, yes!

[aartbik]

addresses comments, no more magic strings

[aartbik]

minor cleanup

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

class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic, [
     SingleBlockImplicitTerminator<"YieldOp">]> {
   let arguments = (ins Variadic<AnyShaped>:$inputs,
                        Variadic<AnyMemRef>:$output_buffers,
                        Variadic<AnyRankedTensor>:$init_tensors,
                        AffineMapArrayAttr:$indexing_maps,
                        ArrayAttr:$iterator_types,
                        OptionalAttr<StrAttr>:$doc,
                        OptionalAttr<StrAttr>:$library_call,
+           // ArrayAttr of StrArrayAttr:
+                       OptionalAttr<ArrayAttr>:$sparse,
                        Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>
                          :$symbol_source);
   let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = [{
     SmallVector<StringRef, 8> linalgTraitAttrNames() {
       return SmallVector<StringRef, 8>{
         getDocAttrName(),
Where #trait_attributes is an alias of a dictionary attribute containing:
         external library function that the linalg.generic operation maps to.
         The external library is assumed to be dynamically linked and no strong
         compile-time guarantees are provided. In the absence of such a library
         call, linalg.generic will always lower to loops.
       - iterator_types: an ArrayAttr specifying the type of the enclosing loops.
         Each element of the list represents and iterator of one of the following
         types:
           parallel, reduction, window
+      - sparse: an optional list with per-dimension sparsity annotations (either
+        "D" for dense or "S" for sparse) for each input and output view.
       - symbol_source: index of the operand whose dimensions will be propagated
         as symbols to the indexing maps. When specified the number of symbols
         in each of the indexing maps has to be either 0 or the rank of the
         specified operand.
 
     Example:
     Defining a #matmul_trait attribute in MLIR can be done as follows:
       ```mlir

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

InterfaceMethod<
           return v.getType().isa<RankedTensorType>();
         };
         return llvm::all_of(getInputs(), isTensorType) &&
                llvm::all_of(this->getOperation()->getResults(), isTensorType);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return whether the op has sparse tensor semantics.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"hasSparseSemantics",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getAttr(getSparseAttrName()).template dyn_cast_or_null<ArrayAttr>() != nullptr;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the name registered for this op when lowering to an external
         library call.
       }],
       /*retTy=*/"std::string",
       /*methodName=*/"getLibraryCallName",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 /// Attribute name for the StringAttr which encodes an optional documentation
 /// string of the structured op.
 constexpr StringRef getDocAttrName() { return "doc"; }
 
 /// Attribute name for the StrArrayAttr which encodes the external library
 /// function that implements the structured op.
 constexpr StringRef getLibraryCallAttrName() { return "library_call"; }
 
+/// Attribute name for the ArrayAttr of StrArrayAttr that encodes sparsity.
+constexpr StringRef getSparseAttrName() { return "sparse"; }
+
 /// Attribute name for the StrArrayAttr which encodes the value of strides.
 constexpr StringRef getStridesAttrName() { return "strides"; }
 
 /// Attribute name for the StrArrayAttr which encodes the value of dilations.
 constexpr StringRef getDilationsAttrName() { return "dilations"; }
 
 /// Attribute name for the StrArrayAttr which encodes the value of paddings.
 constexpr StringRef getPaddingAttrName() { return "padding"; }
inline StringRef toString(IteratorType t) {
   case IteratorType::Parallel:
     return getParallelIteratorTypeName();
   case IteratorType::Reduction:
     return getReductionIteratorTypeName();
   }
   llvm_unreachable("Unsupported IteratorType");
 }
 
+/// Use to encode a dense or sparse dimension.
+constexpr StringRef getSparseDimName() { return "S"; }
+inline bool isSparseDim(Attribute attr) {
+  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
+  return strAttr && strAttr.getValue() == getSparseDimName();
+}
+constexpr StringRef getDenseDimName() { return "D"; }
+inline bool isDenseDim(Attribute attr) {
+  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
+  return strAttr && strAttr.getValue() == getDenseDimName();
+}
+
 } // end namespace mlir
 
 #endif // MLIR_UTILS_STRUCTUREDOPSUTILS_H

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

auto *op =
               types,
               inputValues,
               outputBufferValues,
               initTensorValues,
               builder.getAffineMapArrayAttr(maps),
               builder.getStrArrayAttr(iteratorStrTypes),
               StringAttr() /*doc*/,
               StringAttr() /*library_call*/,
+              ArrayAttr() /*sparse*/,
               IntegerAttr() /*symbol_source*/
               /* TODO: other attributes in op */
               )
           .getOperation();
   // clang-format on
 
   using namespace edsc;
   SmallVector<Type, 4> blockTypes;

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

void GenericOp::build(
     ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorTypes,
     StringRef doc, StringRef libraryCall, IntegerAttr symbolSource,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild) {
   build(builder, result, resultTensorTypes, inputs, outputBuffers, initTensors,
         builder.getAffineMapArrayAttr(indexingMaps),
         builder.getStrArrayAttr(iteratorTypes),
         doc.empty() ? StringAttr() : builder.getStringAttr(doc),
         libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall),
-        symbolSource);
+        ArrayAttr(), symbolSource);
   if (!bodyBuild)
     return;
 
   SmallVector<Type, 4> blockArgTypes;
   for (ValueRange container : {inputs, outputBuffers, initTensors})
     for (Value v : container)
       blockArgTypes.push_back(v.getType().cast<ShapedType>().getElementType());
 
void IndexedGenericOp::build(
     StringRef doc, StringRef libraryCall, IntegerAttr symbolSource,
     function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuild) {
   build(builder, result, resultTensorTypes, inputs, outputBuffers, initTensors,
         builder.getAffineMapArrayAttr(indexingMaps),
         builder.getStrArrayAttr(iteratorTypes),
         doc.empty() ? StringAttr() : builder.getStringAttr(doc),
         libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall),
-        symbolSource);
+        ArrayAttr(), symbolSource);
   if (!bodyBuild)
     return;
 
   unsigned nLoops = iteratorTypes.size();
   SmallVector<Type, 4> blockArgTypes(nLoops, builder.getIndexType());
   for (ValueRange container : {inputs, outputBuffers, initTensors})
     for (Value v : container)
       blockArgTypes.push_back(v.getType().cast<ShapedType>().getElementType());

 void IndexedGenericOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   getGenericEffectsImpl(effects, getOperation()->getResults(),
                         getInputBuffers(), getOutputBuffers());
 }
 
 namespace {
+
 template <typename GenericOpType>
 struct BlockArgsVerifier {
   static LogicalResult verify(GenericOpType op, Block &block);
 };
 
 template <typename GenericOpType>
 LogicalResult BlockArgsVerifier<GenericOpType>::verify(GenericOpType op,
                                                        Block &block) {
if (viewType.getElementType() !=
       return op.emitOpError("expected block argument ")
              << (memrefArgIndex + 1)
              << " of the same type as elemental type of "
              << ((i < nInputViews) ? "input " : "output ")
              << "operand: " << viewType;
   }
   return success();
 }
+
+template <typename GenericOpType>
+struct AnnotationsVerifier {
+  static LogicalResult verify(GenericOpType op) { return success(); }
+};
+
+template <>
+LogicalResult AnnotationsVerifier<GenericOp>::verify(GenericOp op) {
+  ArrayAttr sparseAttr = op.sparseAttr();
+  if (!sparseAttr)
+    return success();
+  // Verify consistency of sparse annotations.
+  if (!op.hasTensorSemantics())
+    return op.emitOpError("expected sparse annotations on tensors only");
+  unsigned numTensors = op.getNumInputsAndOutputs();
+  if (sparseAttr.size() != numTensors)
+    return op.emitOpError("expected one sparse annotation for each tensor");
+  for (unsigned t = 0; t < numTensors; t++) {
+    auto dimAttr = sparseAttr[t].dyn_cast_or_null<ArrayAttr>();
+    if (!dimAttr)
+      return op.emitOpError("expected sparse annotation array for tensor ")
+             << t;
+    unsigned rank = op.getShapedType(t).getRank();
+    if (dimAttr.size() != rank)
+      return op.emitOpError("expected sparse annotation with rank ")
+             << rank << " for tensor " << t;
+    // Per-dimension annotations for each tensor consist of only "D" or "S".
+    for (unsigned d = 0; d < rank; d++) {
+      auto annotation = dimAttr[d].dyn_cast_or_null<StringAttr>();
+      if (isDenseDim(annotation)) {
+        continue;
+      } else if (isSparseDim(annotation)) {
+        if (t == numTensors - 1)
+          return op.emitOpError("sparse output tensors not supported (yet)");
+        continue;
+      }
+      return op.emitOpError("expected sparse annotation at position ")
+             << d << " for tensor " << t;
+    }
+  }
+  return success();
+}
+
 } // namespace
 
 template <typename GenericOpType>
 static LogicalResult verifyGenericOp(GenericOpType op) {
   auto nLoops = op.getNumLoops();
 
   if (op.inputs().size() + op.output_buffers().size() +
           op.init_tensors().size() + op.getNumResults() ==
static LogicalResult verifyGenericOp(GenericOpType op) {
   }
 
   auto concatMap = concatAffineMaps(indexingMaps);
   // TODO: Bound inference for maps with symbols
   if (!concatMap.getNumSymbols() && !inversePermutation(concatMap))
     return op.emitOpError("expected the concatenation of maps in indexing_map "
                           "to be invertible");
 
+  if (failed(AnnotationsVerifier<GenericOpType>::verify(op)))
+    return failure();
+
   return success();
 }
 
 static LogicalResult verify(GenericOp op) { return verifyGenericOp(op); }
 
 static LogicalResult verify(IndexedGenericOp op) { return verifyGenericOp(op); }
 
 //===----------------------------------------------------------------------===//

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

static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
   // Generate a new linalg operation that works on buffers.
   auto newGenericOp = rewriter.create<linalg::GenericOp>(
       genericOp.getLoc(),
       /*resultTensorTypes=*/llvm::None,
       /*inputs=*/inputs,
       /*outputBuffers=*/outputs,
       /*initTensors=*/llvm::None, genericOp.indexing_maps(),
       genericOp.iterator_types(), genericOp.docAttr(),
-      genericOp.library_callAttr(), genericOp.symbol_sourceAttr());
+      genericOp.library_callAttr(), genericOp.sparseAttr(),
+      genericOp.symbol_sourceAttr());
 
   // Create a new block in the region of the new Generic Op.
   Block *oldBlock = genericOp.getBody();
   Region &newRegion = newGenericOp.region();
   Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
                                          oldBlock->getArgumentTypes());
 
   // Add the result arguments to the new block.

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

fusedOp = rewriter
                                      consumer.getOperation()->getResultTypes(),
                                      /*inputs=*/fusedOperands,
                                      /*outputBuffers=*/ValueRange{},
                                      /*initTensors=*/ValueRange{},
                                      rewriter.getArrayAttr(fusedIndexMaps),
                                      consumer.iterator_types(),
                                      /*doc=*/nullptr,
                                      /*library_call=*/nullptr,
+                                     /*sparse=*/nullptr,
                                      /*symbol_source=*/nullptr)
                   .getOperation();
   } else {
     fusedOp =
         rewriter
             .create<IndexedGenericOp>(consumer.getLoc(),
                                       consumer.getOperation()->getResultTypes(),
                                       /*inputs=*/fusedOperands,
                                       /*outputBuffers=*/ValueRange{},
                                       /*initTensors=*/ValueRange{},
                                       rewriter.getArrayAttr(fusedIndexMaps),
                                       consumer.iterator_types(),
                                       /*doc=*/nullptr,
                                       /*library_call=*/nullptr,
+                                      /*sparse=*/nullptr,
                                       /*symbol_source=*/nullptr)
             .getOperation();
   }
 
   // Construct an AffineMap from consumer loops to producer loops.
   // consumer loop -> tensor index
   AffineMap consumerResultIndexMap = consumer.getInputIndexingMap(consumerIdx);
   // producer loop -> tensor index
LinalgOp fusedOp = createLinalgOpOfSameType(
         producer, rewriter, rewriter.getUnknownLoc(), reshapeOp.getResultType(),
         /*inputs=*/producer.getInputs(),
         /*outputBuffers=*/ValueRange{},
         /*initTensors=*/ValueRange{}, // no init tensors for now.
         rewriter.getAffineMapArrayAttr(fusedIndexMaps),
         producer.iterator_types(),
         /*doc=*/nullptr,
         /*library_call=*/nullptr,
+        /*sparse=*/nullptr,
         /*symbol_source=*/nullptr);
     auto &fusedRegion = fusedOp.getOperation()->getRegion(0);
     rewriter.cloneRegionBefore(producer.getOperation()->getRegion(0),
                                fusedRegion, fusedRegion.begin());
     rewriter.replaceOp(reshapeOp, fusedOp.getOperation()->getResults());
     if (producer.use_empty())
       rewriter.eraseOp(producer);
     return success();
for (auto operand : llvm::enumerate(linalgOp.getInputs())) {
           linalgOp.getOperation()->getResultTypes(),
           /*inputs=*/fusedOperands,
           /*outputBuffers=*/ValueRange{},
           /*initTensors=*/ValueRange{}, // no init tensors for now.
           rewriter.getAffineMapArrayAttr(fusedIndexMaps),
           linalgOp.iterator_types(),
           /*doc=*/nullptr,
           /*library_call=*/nullptr,
+          /*sparse=*/nullptr,
           /*symbol_source=*/nullptr);
 
       // Map the block argument corresponding to the replaced argument with the
       // scalar constant.
       Region &linalgOpRegion = linalgOp.getOperation()->getRegion(0);
       Block &entryBlock = *linalgOpRegion.begin();
       unsigned argIndex = entryBlock.getNumArguments() -
                           linalgOp.getNumInputs() + operand.index();

mlir/test/Dialect/Linalg/sparse_invalid.mlir

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+#trait_memref = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  sparse = [
+    [ "S" ],  // a
+    [ "D" ]   // x
+  ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_memref(%arga: memref<32xf32>, %argb: f32) -> tensor<32xf32> {
+ // expected-error@+1 {{'linalg.generic' op expected sparse annotations on tensors only}}
+  %0 = linalg.generic #trait_memref
+    ins(%arga: memref<32xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32xf32>
+  return %0 : tensor<32xf32>
+}
+
+// -----
+
+#trait_too_many = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  sparse = [
+    [ "S" ],  // a
+    [ "S" ],  // b
+    [ "D" ]   // x
+  ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_too_many(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> {
+ // expected-error@+1 {{'linalg.generic' op expected one sparse annotation for each tensor}}
+  %0 = linalg.generic #trait_too_many
+    ins(%arga: tensor<32xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32xf32>
+  return %0 : tensor<32xf32>
+}
+
+// -----
+
+#trait_no_array = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  sparse = [ 1, 2 ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_no_array(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> {
+ // expected-error@+1 {{'linalg.generic' op expected sparse annotation array for tensor 0}}
+  %0 = linalg.generic #trait_no_array
+    ins(%arga: tensor<32xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32xf32>
+  return %0 : tensor<32xf32>
+}
+
+// -----
+
+#trait_wrong_rank = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  sparse = [
+    [ "S" ],
+    [ "D", "D" ]
+  ],
+  iterator_types = ["parallel"]
+}
+
+func @invalid_wrong_rank(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> {
+ // expected-error@+1 {{'linalg.generic' op expected sparse annotation with rank 1 for tensor 1}}
+  %0 = linalg.generic #trait_wrong_rank
+    ins(%arga: tensor<32xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32xf32>
+  return %0 : tensor<32xf32>
+}
+
+// -----
+
+#trait_no_string = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // a
+    affine_map<(i,j) -> (i,j)>   // x (out)
+  ],
+  sparse = [
+    [ "S", 1 ],
+    [ "D", "D" ]
+  ],
+  iterator_types = ["parallel","parallel"]
+}
+
+func @invalid_no_string(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> {
+ // expected-error@+1 {{'linalg.generic' op expected sparse annotation at position 1 for tensor 0}}
+  %0 = linalg.generic #trait_no_string
+    ins(%arga: tensor<32x16xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32x16xf32>
+  return %0 : tensor<32x16xf32>
+}
+
+// -----
+
+#trait_wrong_symbol = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // a
+    affine_map<(i,j) -> (i,j)>   // x (out)
+  ],
+  sparse = [
+    [ "S", "S" ],
+    [ "D", "X" ]
+  ],
+  iterator_types = ["parallel","parallel"]
+}
+
+func @invalid_wrong_symbol(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> {
+  // expected-error@+1 {{'linalg.generic' op expected sparse annotation at position 1 for tensor 1}}
+  %0 = linalg.generic #trait_wrong_symbol
+    ins(%arga: tensor<32x16xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32x16xf32>
+  return %0 : tensor<32x16xf32>
+}
+
+// -----
+
+#trait_no_sparse_output = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // a
+    affine_map<(i,j) -> (i,j)>   // x (out)
+  ],
+  sparse = [
+    [ "S", "S" ],
+    [ "D", "S" ]
+  ],
+  iterator_types = ["parallel","parallel"]
+}
+
+func @invalid_no_sparse_output(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> {
+  // expected-error@+1 {{'linalg.generic' op sparse output tensors not supported (yet)}}
+  %0 = linalg.generic #trait_no_sparse_output
+    ins(%arga: tensor<32x16xf32>) {
+      ^bb(%a: f32):
+        %0 = addf %a, %argb  : f32
+        linalg.yield %0 : f32
+  } -> tensor<32x16xf32>
+  return %0 : tensor<32x16xf32>
+}