[mlir][linalg] remove the -now- obsolete sparse support in linalg

[aartbik]

Summary

All glue and clutter in the linalg ops has been replaced by proper
sparse tensor type encoding. This code is no longer needed. Thanks
to ntv@ for giving us a temporary home in linalg.

So long, and thanks for all the fish.

[aartbik]

rebased with parent revision : ready to go now!

[aartbik]

rebase even further

[aartbik]

rebase

[aartbik]

rebase even more, trying to get passed recent webasm breakages...

[aartbik]

rebase past the webasm revert

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

InterfaceMethod<
       /*defaultImplementation=*/[{
         return llvm::all_of(getShapedOperands(), [](Value v) {
           return v.getType().template isa<RankedTensorType>();
         });
       }]
     >,
     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/Linalg/IR/LinalgStructuredOps.td

class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic, [
     AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
     SingleBlockImplicitTerminator<"YieldOp">]> {
   let arguments = (ins Variadic<AnyShaped>:$inputs,
                        Variadic<AnyShaped>:$outputs,
                        AffineMapArrayAttr:$indexing_maps,
                        ArrayAttr:$iterator_types,
                        OptionalAttr<StrAttr>:$doc,
-                       OptionalAttr<StrAttr>:$library_call,
-                       // ArrayAttr of StrArrayAttr:
-                       OptionalAttr<ArrayAttr>:$sparse);
+                       OptionalAttr<StrAttr>:$library_call);
   let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = structuredOpsBaseDecls # [{
     SmallVector<StringRef, 8> linalgTraitAttrNames() {
       return SmallVector<StringRef, 8>{
         getDocAttrName(),
         getIndexingMapsAttrName(), getLibraryCallAttrName(),
         getIteratorTypesAttrName(),
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.
 
     Example:
     Defining a #matmul_trait attribute in MLIR can be done as follows:
       ```mlir
       #matmul_accesses = [
         (m, n, k) -> (m, k),
         (m, n, k) -> (k, n),
         (m, n, k) -> (m, n)

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/IR/LinalgOps.cpp

void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild) {
   build(builder, result, resultTensorTypes, inputs, outputs,
         builder.getAffineMapArrayAttr(indexingMaps),
         builder.getStrArrayAttr(iteratorTypes),
         doc.empty() ? StringAttr() : builder.getStringAttr(doc),
-        libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall),
-        ArrayAttr());
+        libraryCall.empty() ? StringAttr()
+                            : builder.getStringAttr(libraryCall));
   if (!bodyBuild)
     return;
 
   SmallVector<Type, 4> blockArgTypes;
   for (ValueRange container : {inputs, outputs})
     for (Value v : container)
       blockArgTypes.push_back(v.getType().cast<ShapedType>().getElementType());
 
void IndexedGenericOp::build(
     ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuild) {
   build(builder, result, resultTensorTypes, inputs, outputs,
         builder.getAffineMapArrayAttr(indexingMaps),
         builder.getStrArrayAttr(iteratorTypes),
         doc.empty() ? StringAttr() : builder.getStringAttr(doc),
-        libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall),
-        ArrayAttr());
+        libraryCall.empty() ? StringAttr()
+                            : builder.getStringAttr(libraryCall));
   if (!bodyBuild)
     return;
 
   unsigned nLoops = iteratorTypes.size();
   SmallVector<Type, 4> blockArgTypes(nLoops, builder.getIndexType());
   for (ValueRange container : {inputs, outputs})
     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 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");
-  if (op.getNumOutputs() != 1)
-    return op.emitOpError("expected single output tensor");
-  unsigned numTensors = op.getNumShapedOperands();
-  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++) {
-      if (isDenseDim(dimAttr[d])) {
-        continue;
-      } else if (isSparseDim(dimAttr[d])) {
-        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) {
-  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); }
 
 namespace {

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

finalizeBufferAllocationForGenericOp(ConversionPatternRewriter &rewriter,
                                      ValueRange outputs) {
   // Generate a new linalg operation that works on buffers.
   auto newGenericOp = rewriter.create<GenericOpTy>(
       genericOp.getLoc(),
       /*resultTensorTypes=*/llvm::None,
       /*inputs=*/inputs,
       /*outputs=*/outputs, genericOp.indexing_maps(),
       genericOp.iterator_types(), genericOp.docAttr(),
-      genericOp.library_callAttr(), genericOp.sparseAttr());
+      genericOp.library_callAttr());
 
   // 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());
 
   // Clone the body of the old block to the new block.

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

if (isa<GenericOp>(producer.getOperation()) &&
       isa<GenericOp>(consumer.getOperation())) {
     fusedOp = rewriter.create<GenericOp>(
         consumer.getLoc(), consumer->getResultTypes(),
         /*inputs=*/fusedOperands,
         // TODO: handle outputs.
         consumer.getOutputs(), rewriter.getAffineMapArrayAttr(fusedIndexMaps),
         consumer.iterator_types(),
         /*doc=*/nullptr,
-        /*library_call=*/nullptr,
-        /*sparse=*/nullptr);
+        /*library_call=*/nullptr);
   } else {
     fusedOp = rewriter.create<IndexedGenericOp>(
         consumer.getLoc(), consumer->getResultTypes(),
         /*inputs=*/fusedOperands,
         // TODO: handle outputs.
         consumer.getOutputs(), rewriter.getAffineMapArrayAttr(fusedIndexMaps),
         consumer.iterator_types(),
         /*doc=*/nullptr,
-        /*library_call=*/nullptr,
-        /*sparse=*/nullptr);
+        /*library_call=*/nullptr);
   }
 
   // Construct an AffineMap from consumer loops to producer loops.
   // consumer loop -> tensor index
   AffineMap consumerResultIndexMap = consumer.getInputIndexingMap(consumerIdx);
   // tensor index -> producer loop
   AffineMap invProducerResultIndexMap =
       inversePermutation(producerResultIndexMap);
Value output = rewriter.create<TensorReshapeOp>(
         loc, producer.getOutputs()[0], reshapeOp.getReassociationExprs());
     LinalgOp fusedOp = createLinalgOpOfSameType(
         producer, rewriter, loc, reshapeOp.getResultType(),
         /*inputs=*/producer.getInputs(),
         // TODO: handle outputs.
         /*outputs=*/output, rewriter.getAffineMapArrayAttr(fusedIndexMaps),
         producer.iterator_types(),
         /*doc=*/nullptr,
-        /*library_call=*/nullptr,
-        /*sparse=*/nullptr);
+        /*library_call=*/nullptr);
     auto &fusedRegion = fusedOp->getRegion(0);
     rewriter.cloneRegionBefore(producer->getRegion(0), fusedRegion,
                                fusedRegion.begin());
     rewriter.replaceOp(reshapeOp, fusedOp->getResults());
     return success();
   }
 };
 
for (auto operand : llvm::enumerate(linalgOp.getInputOpOperands())) {
       LinalgOp fusedOp = createLinalgOpOfSameType(
           linalgOp, rewriter, rewriter.getUnknownLoc(),
           linalgOp->getResultTypes(),
           /*inputs=*/fusedOperands,
           /*outputs=*/linalgOp.getOutputs(),
           rewriter.getAffineMapArrayAttr(fusedIndexMaps),
           linalgOp.iterator_types(),
           /*doc=*/nullptr,
-          /*library_call=*/nullptr,
-          /*sparse=*/nullptr);
+          /*library_call=*/nullptr);
 
       // Map the block argument corresponding to the replaced argument with the
       // scalar constant.
       Region &linalgOpRegion = linalgOp->getRegion(0);
       Block &entryBlock = *linalgOpRegion.begin();
       unsigned argIndex = entryBlock.getNumArguments() -
                           linalgOp.getNumShapedOperands() + operand.index();
       BlockAndValueMapping mapping;

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
   # TODO: Support emission of pure memref form.
   indexing_maps_attr = ArrayAttr.get(
       [AffineMapAttr.get(am)
        # TODO: linalg verification does not currently allow symbols.
        # Compress them for now.
        for am in AffineMap.compress_unused_symbols(op_config.indexing_maps, Context.current)])
   iterator_types_attr = ArrayAttr.get(
       [StringAttr.get(s) for s in op_config.iterator_types])
-  # TODO: Add support for sparse operands once there is a stable interface.
-  sparse_attr = None
 
   return (all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types,
           type_mapping, capture_arg_mapping, indexing_maps_attr,
-          iterator_types_attr, sparse_attr)
+          iterator_types_attr)
 
 
 def emit_generic_structured_op(op_config: LinalgStructuredOpConfig,
                                *ins: Value,
                                outs: Sequence[Value] = (),
                                captures: Sequence[Value] = ()):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
-  capture_arg_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr = \
+  capture_arg_mapping, indexing_maps_attr, iterator_types_attr = \
      prepare_common_structured_op(op_config, *ins, outs = outs,
                                   captures=captures)
 
   generic_op = linalg.GenericOp(
       result_tensors=result_types,
       inputs=ins,
       outputs=outs,
       indexing_maps=indexing_maps_attr,
       iterator_types=iterator_types_attr,
       doc=None,  # TODO: Make optional.
-      library_call=None,  # TODO: Make optional.
-      sparse=sparse_attr)  # TODO: Make optional.
+      library_call=None)  # TODO: Make optional.
 
   # Construct the body.
   block_arg_names = _get_tensor_def_names(*in_arg_defs, *out_arg_defs)
   block_arg_types = _get_shaped_element_types_from_values(*ins, *outs)
   block = generic_op.regions[0].blocks.append(*block_arg_types)
   block_arg_mapping = dict(zip(block_arg_names, block.arguments))
   with InsertionPoint(block):
     body_builder = _BodyBuilder(type_mapping, block_arg_mapping,

 
 def emit_named_structured_op(op_config: LinalgStructuredOpConfig,
                              op_name: str,
                              op_class_name: str,
                              *ins: Value,
                              outs: Sequence[Value] = (),
                              captures: Sequence[Value] = ()):
   all_arg_defs, in_arg_defs, out_arg_defs, outs, result_types, type_mapping, \
-  capture_arg_mapping, indexing_maps_attr, iterator_types_attr, sparse_attr = \
+  capture_arg_mapping, indexing_maps_attr, iterator_types_attr = \
      prepare_common_structured_op(op_config, *ins, outs = outs,
                                   captures = captures)
 
   # If we get here, there must exist a builtin class `op_class_name`.
   ctx = Context.current
   fully_qualified_name = 'linalg.' + op_name
   if (not ctx.is_registered_operation(fully_qualified_name) or
       not op_class_name in linalg.__dict__.keys()):

 
 def _get_tensor_def_names(
     *tensor_def_configs: TensorDefConfig) -> Sequence[str]:
   return [tdc.tensor_def.tensor_name for tdc in tensor_def_configs]
 
 def _add_type_mapping(name: str, type: Type, type_mapping: Dict[str, Type]):
   if name in type_mapping:
     if type_mapping[name] != type:
-        raise ValueError(f"Cannot overwrite type mapping {name} = "
-                         f"{type_mapping[name]} by type {type}")
   type_mapping[name] = type
 
 def _is_floating_point_type(t: Type) -> bool:
   # TODO: Create a FloatType in the Python API and implement the switch
   # there.
   return (F64Type.isinstance(t) or F32Type.isinstance(t) or
           F16Type.isinstance(t) or BF16Type.isinstance(t))