[mlir][WIP] symbol_source attr in linalg.generic proposal

[limo1996]

Summary

Depends On D83191

[limo1996]

Linting message fixed

[ftynse]

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

522

I think llvm::Optional<int> is preferable to returning implicit -1 as a sign of the attribute being absent.

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

321

Do you see a good reason why this code (and below) is using the default autogenerated builder form instead of the custom builder form? Maybe it is worth updating so that next modifications of the Op arguments don't require changes.

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

529

We shouldn't be having this attribute on the IndexedGenericOp.

mlir/test/Dialect/Linalg/invalid.mlir

124

There's no need to check autogenerated error messages

[limo1996]

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

522

Ahh Thank you. Didn't know about this option.

mlir/test/Dialect/Linalg/invalid.mlir

124

Good to know. Should I remove the check completely?

[limo1996]

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

529

Yeah that's kind of the point of this revision.. I know we should not but I wanted to ask how can I solve it bcos I don't know how to add to arguments instead of rewrite them inside GenericOp definition.. Any suggestion is rly welcomed :)

[ftynse]

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

522

Just change the return type and return llvm::None instead of -1.

mlir/test/Dialect/Linalg/invalid.mlir

124

Yes

[limo1996]

getSymbolSource: llvm::Optional instead of -1

[limo1996]

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

321

So the problem is that custom builder requires types like int64_t etc. but autogenerated one requires corresponding attributes.. So I would suggest one more custom builder but with attributes. What do you think?

[limo1996]

Test cases for now autogenerated error messages removed

[limo1996]

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

321

Submited it as a bug. You can find it here: https://bugs.llvm.org/show_bug.cgi?id=46658

[limo1996]

changes evolved from revision D83158

[limo1996]

evolve

[limo1996]

Changes merged to D83158

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

 class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic,
     [SingleBlockImplicitTerminator<"YieldOp">]> {
   let arguments = (ins Variadic<LinalgOperand>:$views,
                    I64Attr:$args_in,
                    I64Attr:$args_out,
                    AffineMapArrayAttr:$indexing_maps,
                    ArrayAttr:$iterator_types,
                    OptionalAttr<StrAttr>:$doc,
-                   OptionalAttr<StrAttr>:$library_call);
+                   OptionalAttr<StrAttr>:$library_call,
+                   Confined<OptionalAttr<I64Attr>,
+                     [IntMinValue<0>]>:$symbol_source);
   let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = [{
     SmallVector<StringRef, 8> linalgTraitAttrNames() {
       return SmallVector<StringRef, 8>{
         getArgsInAttrName(), getArgsOutAttrName(), getDocAttrName(),
         getIndexingMapsAttrName(), getLibraryCallAttrName(),
-        getIteratorTypesAttrName()
+        getIteratorTypesAttrName(), getSymbolSourceAttrName()
       };
     }
 
     unsigned getNumInputs() { return args_in().getSExtValue(); }
 
     unsigned getNumOutputs() { return args_out().getSExtValue(); }
 
     StringRef getLibraryCallName() {
       return library_call().hasValue() ? library_call().getValue() : "";
     }
 
     llvm::Optional<SmallVector<StringRef, 8>> referenceIterators() {
       llvm_unreachable(
         "No such thing as reference iterator types for a generic op.");
      }
 
     llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
       llvm_unreachable(
         "No such thing as reference indexing maps for a generic op.");
      }
+
+     int getSymbolSource() {
+       auto ss = symbol_source();
+       return ss.hasValue() ? ss.getValue().getLimitedValue() : -1;
+     }
   }];
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parseGenericOp(parser, result); }];
 }
 
 /// Index-free GenericOp.
 def GenericOp : GenericOpBase<"generic"> {

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

inline bool isColumnMajorMatmul(ArrayAttr indexingMaps) {
   bindDims(context, m, n, k);
   auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
   auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
   auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
   auto maps = ArrayAttr::get({mapA, mapB, mapC}, context);
   return indexingMaps == maps;
 }
 
+/// Attribute name for the IntegerAttr which encodes the index of operand
+/// whose dimensions will be propagated as symbols to the indexing maps
+constexpr StringRef getSymbolSourceAttrName() { return "symbol_source"; }
+
 /// Attribute name for the AffineArrayAttr which encodes the relationship
 /// between a structured op iterators' and its operands.
 constexpr StringRef getIndexingMapsAttrName() { return "indexing_maps"; }
 
 /// Attribute name for the StrArrayAttr which encodes the type of a structured
 /// op's iterators.
 constexpr StringRef getIteratorTypesAttrName() { return "iterator_types"; }
 

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

auto *op =
               edsc::ScopedContext::getLocation(),
               types,
               values,
               IntegerAttr::get(IntegerType::get(64, ctx), nInputs),
               IntegerAttr::get(IntegerType::get(64, ctx), nOutputs),
               builder.getAffineMapArrayAttr(maps),
               builder.getStrArrayAttr(iteratorStrTypes),
               StringAttr() /*doc*/,
-              StringAttr() /*library_call*/
+              StringAttr() /*library_call*/,
+              IntegerAttr() /*symbol_source*/
               /* TODO: other attributes in op */
               )
           .getOperation();
   // clang-format on
 
   using namespace edsc;
   SmallVector<Type, 4> blockTypes;
   blockTypes.reserve(values.size());

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

void GenericOp::build(
     ValueRange args, int64_t inputCount, int64_t outputCount,
     ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorTypes,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild) {
   build(builder, result, resultTypes, args,
         builder.getI64IntegerAttr(inputCount),
         builder.getI64IntegerAttr(outputCount),
         builder.getAffineMapArrayAttr(indexingMaps),
         builder.getStrArrayAttr(iteratorTypes),
-        /*doc=*/nullptr, /*library_call=*/nullptr);
+        /*doc=*/nullptr, /*library_call=*/nullptr,
+        /*symbol_source=*/nullptr);
   if (!bodyBuild)
     return;
 
   SmallVector<Type, 4> blockArgTypes;
   for (Value arg : args)
     blockArgTypes.push_back(arg.getType().cast<ShapedType>().getElementType());
 
   OpBuilder::InsertionGuard guard(builder);
   auto &region = *result.regions.front();
   Block *bodyBlock = builder.createBlock(&region, region.end(), blockArgTypes);
   bodyBuild(builder, result.location, bodyBlock->getArguments());
 }
 
 void IndexedGenericOp::build(
     OpBuilder &builder, OperationState &result, ArrayRef<Type> resultTypes,
     ValueRange args, int64_t inputCount, int64_t outputCount,
     ArrayRef<AffineMap> indexingMaps, ArrayRef<StringRef> iteratorTypes,
     function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>
         bodyBuild) {
   build(builder, result, resultTypes, args,
         builder.getI64IntegerAttr(inputCount),
         builder.getI64IntegerAttr(outputCount),
         builder.getAffineMapArrayAttr(indexingMaps),
         builder.getStrArrayAttr(iteratorTypes),
-        /*doc=*/nullptr, /*library_call=*/nullptr);
+        /*doc=*/nullptr, /*library_call=*/nullptr,
+        /*symbol_source=*/nullptr);
   if (!bodyBuild)
     return;
 
   unsigned nLoops = iteratorTypes.size();
   SmallVector<Type, 4> blockArgTypes(nLoops, builder.getIndexType());
   for (Value arg : args)
     blockArgTypes.push_back(arg.getType().cast<ShapedType>().getElementType());
 
return op.emitOpError("expected exactly ")
            << " inputs (tensor or buffer) and output buffer operands";
 
   auto &region = op.region();
   if (!llvm::hasSingleElement(region))
     return op.emitOpError("expected region with 1 block");
   if (failed(BlockArgsVerifier<GenericOpType>::verify(op, region.front())))
     return failure();
 
-  auto attr = op.getAttr("symbol_source");
+  auto attr = op.template getAttrOfType<IntegerAttr>("symbol_source");
   int64_t targetRank = 0;
   if (attr) {
-    if (attr.getKind() != StandardAttributes::Kind::Integer)
-      return op.emitOpError("symbol_source attribute has to be integer");
-    auto index = attr.template cast<IntegerAttr>().getInt();
-    if (index < 0 || index >= op.getNumOperands())
+    unsigned index = attr.getInt();
+    if (index >= op.getNumOperands())
       return op.emitOpError("symbol_source index out of range");
     targetRank =
         op.getOperand(index).getType().template cast<ShapedType>().getRank();
   }
 
   SmallVector<AffineMap, 4> indexingMaps;
   indexingMaps.reserve(op.indexing_maps().size());
   for (auto en : llvm::enumerate(op.indexing_maps())) {

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

LogicalResult matchAndRewrite(GenericOp genericOp,
     resultTypes.reserve(genericOp.getNumResults());
     for (unsigned i : llvm::seq<unsigned>(0, genericOp.getNumResults()))
       resultTypes.push_back(
           newInputOutputTypes[i + genericOp.getNumOperands()]);
     GenericOp replacementOp = rewriter.create<GenericOp>(
         loc, resultTypes, newOperands, genericOp.args_in(),
         genericOp.args_out(), rewriter.getAffineMapArrayAttr(newIndexingMaps),
         genericOp.iterator_types(),
         /*doc = */ nullptr,
-        /*library_call = */ nullptr);
+        /*library_call = */ nullptr,
+        /*symbol_source = */ nullptr);
     rewriter.inlineRegionBefore(genericOp.region(), replacementOp.region(),
                                 replacementOp.region().begin());
 
     // If any result tensor has a modified shape, then add reshape to recover
     // the original shape.
     SmallVector<Value, 4> resultReplacements;
     for (auto result : llvm::enumerate(replacementOp.getResults())) {
       unsigned index = result.index() + replacementOp.getNumOperands();

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

if (isa<GenericOp>(producer.getOperation()) &&
                   rewriter.getUnknownLoc(),
                   consumer.getOperation()->getResultTypes(), fusedOperands,
                   rewriter.getI64IntegerAttr(fusedOperands.size()),
                   rewriter.getI64IntegerAttr(
                       consumer.getOperation()->getNumResults()),
                   rewriter.getArrayAttr(fusedIndexMaps),
                   consumer.iterator_types(),
                   /*doc=*/nullptr,
-                  /*library_call=*/nullptr)
+                  /*library_call=*/nullptr,
+                  /*symbol_source=*/nullptr)
               .getOperation();
     } else {
       fusedOp =
           rewriter
               .create<IndexedGenericOp>(
                   rewriter.getUnknownLoc(),
                   consumer.getOperation()->getResultTypes(), fusedOperands,
                   rewriter.getI64IntegerAttr(fusedOperands.size()),
                   rewriter.getI64IntegerAttr(
                       consumer.getOperation()->getNumResults()),
                   rewriter.getArrayAttr(fusedIndexMaps),
                   consumer.iterator_types(),
                   /*doc=*/nullptr,
-                  /*library_call=*/nullptr)
+                  /*library_call=*/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
SmallVector<Attribute, 4> indexMapAttrs = llvm::to_vector<4>(
           return AffineMapAttr::get(map);
         }));
     auto fusedOp = rewriter.create<LinalgOpTy>(
         rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands,
         rewriter.getI64IntegerAttr(fusedOperands.size()),
         rewriter.getI64IntegerAttr(consumer.getNumResults()),
         rewriter.getArrayAttr(indexMapAttrs), consumer.iterator_types(),
         /*doc=*/nullptr,
-        /*library_call=*/nullptr);
+        /*library_call=*/nullptr,
+        /*symbol_source=*/nullptr);
     auto &fusedRegion = fusedOp.region();
     rewriter.cloneRegionBefore(consumer.region(), fusedRegion,
                                fusedRegion.begin());
     return fusedOp;
   }
 };
 
 /// Implementation of fusion on tensor ops when consumer is a TensorReshapeOp.
static Operation *fuse(LinalgOpTy producer, TensorReshapeOp consumer,
 
     auto fusedOp = rewriter.create<LinalgOpTy>(
         rewriter.getUnknownLoc(), consumer.getResultType(),
         producer.getOperands(),
         rewriter.getI64IntegerAttr(producer.getNumOperands()),
         rewriter.getI64IntegerAttr(1), rewriter.getArrayAttr(indexMapAttrs),
         producer.iterator_types(),
         /*doc=*/nullptr,
-        /*library_call=*/nullptr);
+        /*library_call=*/nullptr,
+        /*symbol_source=*/nullptr);
     auto &fusedRegion = fusedOp.region();
     rewriter.cloneRegionBefore(producer.region(), fusedRegion,
                                fusedRegion.begin());
     return fusedOp;
   }
 };
 
 /// Implementation of fusion on tensor ops when producer is a splat constant.
static Operation *fuse(ConstantOp producer, LinalgOpTy consumer,
 
     auto fusedOp = rewriter.create<LinalgOpTy>(
         rewriter.getUnknownLoc(), consumer.getResultTypes(), fusedOperands,
         rewriter.getI64IntegerAttr(consumer.getNumOperands() - 1),
         rewriter.getI64IntegerAttr(consumer.getNumResults()),
         rewriter.getAffineMapArrayAttr(fusedIndexMaps),
         consumer.iterator_types(),
         /*doc=*/nullptr,
-        /*library_call=*/nullptr);
+        /*library_call=*/nullptr,
+        /*symbol_source=*/nullptr);
 
     // Map the block argument corresponding to the replaced argument with the
     // scalar constant.
     Region &consumerRegion = consumer.region();
     Block &entryBlock = *consumerRegion.begin();
     unsigned argIndex =
         entryBlock.getNumArguments() - consumer.getNumOperands() + consumerIdx;
     BlockAndValueMapping mapping;

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

for (auto result : results) {
       newArgs.push_back(alloc);
       newResults.push_back(alloc);
     }
 
     // Generate a new linalg operation that works on buffers.
     auto linalgOp = rewriter.create<linalg::GenericOp>(
         loc, llvm::None, newArgs, rewriter.getI64IntegerAttr(operands.size()),
         rewriter.getI64IntegerAttr(results.size()), op.indexing_maps(),
-        op.iterator_types(), op.docAttr(), op.library_callAttr());
+        op.iterator_types(), op.docAttr(), op.library_callAttr(),
+        op.symbol_sourceAttr());
 
     // Create a new block in the region of the new Generic Op.
     Block &oldBlock = op.getRegion().front();
     Region &newRegion = linalgOp.region();
     Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
                                            oldBlock.getArgumentTypes());
 
     // Add the result arguments to the new block.

mlir/test/Dialect/Linalg/invalid.mlir

linalg.generic {
     ^bb(%i : i32):
     linalg.yield %i : i32
   }: memref<i32>
 }
 
 // -----
 
 func @generic_symbol_source_wrong_type(%arg0: memref<i32>) {
-  // expected-error @+1 {{symbol_source attribute has to be integer}}
+  // expected-error @+1 {{linalg.generic' op attribute 'symbol_source' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<()[N] -> (0)> ],
     iterator_types = ["parallel"],
     symbol_source = "none"
   } %arg0 {
     ^bb(%i : i32):
linalg.generic {
     ^bb(%i : i32):
     linalg.yield %i : i32
   }: memref<i32>
 }
 
 // -----
 
 func @generic_symbol_source_out_of_range(%arg0: memref<i32>) {
-  // expected-error @+1 {{symbol_source index out of range}}
+  // expected-error @+1 {{linalg.generic' op attribute 'symbol_source' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<()[N] -> (0)> ],
     iterator_types = ["parallel"],
     symbol_source = -1
   } %arg0 {
     ^bb(%i : i32):

mlir/test/lib/Transforms/TestBufferPlacement.cpp

matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
         newArgs.push_back(alloc);
         newResults.push_back(alloc);
       }
 
       // Generate a new linalg operation that works on buffers.
       auto linalgOp = rewriter.create<linalg::GenericOp>(
           loc, llvm::None, newArgs, rewriter.getI64IntegerAttr(operands.size()),
           rewriter.getI64IntegerAttr(results.size()), op.indexing_maps(),
-          op.iterator_types(), op.docAttr(), op.library_callAttr());
+          op.iterator_types(), op.docAttr(), op.library_callAttr(),
+          op.symbol_sourceAttr());
 
       // Create a new block in the region of the new Generic Op.
       Block &oldBlock = op.getRegion().front();
       Region &newRegion = linalgOp.region();
       Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
                                              oldBlock.getArgumentTypes());
 
       // Map the old block arguments to the new ones.