[mlir] Remove iterator_types() method from LinalgStructuredInterface.

[olegshyshkov]

Summary

getIteratorTypesArray should be used instead. It's a better substitute for all the current usages of the interface.

The current ArrayAttr iterator_types() has a few problems:

• It creates an assumption operation has iterators types as an attribute, but it's not always the case. Sometime iterator types can be inferred from other attribute, or they're just static.
• ArrayAttr is an obscure contained and required extracting values in the client code.
• Makes it hard to migrate iterator types from strings to enums (RFC).

Concrete ops, like linalg.generic will still have iterator types as an attribute if needed.

As a side effect, this change helps a bit with migration to prefixed accessors.

[olegshyshkov]

Typos.

[pifon2a]

Nice!

[nicolasvasilache]

thanks!

[olegshyshkov]

Sync and resolve merge conflicts with Linalg::MapOp that was just added in https://reviews.llvm.org/D135754.

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

InterfaceMethod<
       /*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.getIteratorTypes();
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/[{
         Return iterator types in the current operation.
+
+        Default implementation assumes that the operation has an attribute
+        `iterator_types`, but it's not always the case. Sometimes iterator types
+        can be infered from other parameters and in such cases default
+        getIteratorTypesArray should be overriden.
       }],
       /*retTy=*/"SmallVector<StringRef>",
       /*methodName=*/"getIteratorTypesArray",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        auto range = $_op.iterator_types().template getAsValueRange<StringAttr>();
+        auto range = $_op.getIteratorTypes().template getAsValueRange<StringAttr>();
         return {range.begin(), range.end()};
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return true if the indexing map is depending on the current op instance.
         This means that the indexing map is dynamically synthesized by using the
         op instance's concrete attributes, instead of being static for all
let extraClassDeclaration = [{
     /// `createFlatListOfOperandStaticDims`.
     SmallVector<int64_t, 4> computeStaticLoopSizes();
 
     /// Returns the value that expresses the shape of the output in terms of
     /// shape of the input operands where possible
     LogicalResult reifyResultShapes(OpBuilder &b,
         ReifiedRankedShapedTypeDims &reifiedReturnShapes);
 
-    // TODO: Remove once prefixing is flipped.
-    ArrayAttr getIteratorTypes() { return iterator_types(); }
-
     SmallVector<StringRef> getIteratorTypeNames() {
       return getIteratorTypesArray();
     }
 
     //========================================================================//
     // Forwarding functions to access interface methods from the
     // DestinationStyleOpInterface.
     // MLIR currently does not support dependent interfaces or interface

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

let arguments = (ins
     // Output arg
     TensorOrMemref:$init
   );
   let results = (outs Variadic<AnyTensor>:$result);
   let regions = (region SizedRegion<1>:$mapper);
 
   let extraClassDeclaration = structuredOpsBaseDecls # [{
     // Implement functions necessary for LinalgStructuredInterface.
-    ArrayAttr getIteratorTypes();
+    SmallVector<StringRef> getIteratorTypesArray();
     ArrayAttr getIndexingMaps();
     std::string getLibraryCallName() {
       return "op_has_no_registered_library_name";
     }
 
     // Implement functions necessary for DestinationStyleOpInterface.
     unsigned getNumInputs() {
       return this->getOperation()->getNumOperands() - getNumOutputs();
let arguments = (ins
     ConfinedAttr<DenseI64ArrayAttr,
                  [DenseArrayStrictlySorted<DenseI64ArrayAttr>]>:$dimensions
   );
   let results = (outs Variadic<AnyTensor>);
   let regions = (region SizedRegion<1>:$combiner);
 
   let extraClassDeclaration = structuredOpsBaseDecls # [{
     // Declare functions necessary for LinalgStructuredInterface.
-    ArrayAttr getIteratorTypes();
+    SmallVector<StringRef> getIteratorTypesArray();
     ArrayAttr getIndexingMaps();
     std::string getLibraryCallName() {
       return "op_has_no_registered_library_name";
     }
 
     // Implement functions necessary for DestinationStyleOpInterface.
     mlir::ValueRange getOutputs() { return getInits(); }
     unsigned getNumInputs() { return getInputs().size(); };

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

if (inputElemShape != outputShape) {
                            << ") to match shape of output (" << outputShape
                            << ")";
     }
   }
 
   return success();
 }
 
-ArrayAttr MapOp::getIteratorTypes() {
+SmallVector<StringRef> MapOp::getIteratorTypesArray() {
   int64_t rank = getInit().getType().getRank();
-  return Builder(getContext())
-      .getStrArrayAttr(
-          SmallVector<StringRef>(rank, getParallelIteratorTypeName()));
+  return SmallVector<StringRef>(rank, getParallelIteratorTypeName());
 }
 
 ArrayAttr MapOp::getIndexingMaps() {
   Builder builder(getContext());
   int64_t rank = getInit().getType().getRank();
   int64_t numIndexingMaps = getOperands().size();
   return builder.getAffineMapArrayAttr(SmallVector<AffineMap>(
       numIndexingMaps, builder.getMultiDimIdentityMap(rank)));

 }
 
 void ReduceOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "reduced");
 }
 
-ArrayAttr ReduceOp::getIteratorTypes() {
+SmallVector<StringRef> ReduceOp::getIteratorTypesArray() {
   int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
   SmallVector<StringRef> iteratorTypes(inputRank,
                                        getParallelIteratorTypeName());
   for (int64_t reductionDim : getDimensions())
     iteratorTypes[reductionDim] = getReductionIteratorTypeName();
-  return Builder(getContext()).getStrArrayAttr(iteratorTypes);
+  return iteratorTypes;
 }
 
 ArrayAttr ReduceOp::getIndexingMaps() {
   int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
   SmallVector<AffineMap> affineMaps(
       getNumInputs(),
       AffineMap::getMultiDimIdentityMap(inputRank, getContext()));
   AffineMap resultMap =

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

struct LinalgOpTilingInterface
   /// Return the destination operands.
   SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
     return cast<LinalgOp>(op).getOutputOperands();
   }
 
   /// Return the loop iterator type.
   SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
     LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
-    return llvm::to_vector(
-        llvm::map_range(concreteOp.iterator_types(), [](Attribute strAttr) {
-          return utils::symbolizeIteratorType(
-                     strAttr.cast<StringAttr>().getValue())
-              .value();
+    return llvm::to_vector(llvm::map_range(
+        concreteOp.getIteratorTypesArray(), [](StringRef iteratorType) {
+          return utils::symbolizeIteratorType(iteratorType).value();
         }));
   }
 
   /// Return the iteration domain range.
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
     Location loc = op->getLoc();

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

LogicalResult matchAndRewrite(GenericOp op,
     SmallVector<Value> inputOps = prod.getInputOperands();
     SmallVector<Value> outputOps = op.getOutputOperands();
     SmallVector<AffineMap> fusedIndexMaps = prod.getIndexingMapsArray();
     inputOps.push_back(op.getInputOperand(1 - other)->get());
     fusedIndexMaps.push_back(fusedIndexMaps.back()); // mimic other
     // Fuse producer and consumer into a new generic op.
     auto fusedOp = rewriter.create<GenericOp>(
         loc, op.getResult(0).getType(), inputOps, outputOps,
-        rewriter.getAffineMapArrayAttr(fusedIndexMaps), prod.iterator_types(),
+        rewriter.getAffineMapArrayAttr(fusedIndexMaps), prod.getIteratorTypes(),
         /*doc=*/nullptr, /*library_call=*/nullptr);
     Block &prodBlock = prod.getRegion().front();
     Block &consBlock = op.getRegion().front();
     BlockAndValueMapping mapper;
     Block *fusedBlock = new Block();
     fusedOp.getRegion().push_back(fusedBlock);
     unsigned num = prodBlock.getNumArguments();
     for (unsigned i = 0; i < num - 1; i++)

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

public:
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Detects sparse annotations and translate the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     if (op.getNumOutputs() != 1)
       return failure();
     unsigned numTensors = op.getNumInputsAndOutputs();
-    unsigned numLoops = op.iterator_types().getValue().size();
+    unsigned numLoops = op.getNumLoops();
     Merger merger(numTensors, numLoops);
     if (!findSparseAnnotations(merger, op))
       return failure();
 
     // Builds the tensor expression for the Linalg operation in SSA form.
     Optional<unsigned> optExp = merger.buildTensorExpFromLinalg(op);
     if (!optExp.has_value())
       return failure();

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

let extraClassDeclaration = [{
     }
 
     static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
                               mlir::ArrayRef<mlir::NamedAttribute>)>
     getRegionBuilder() {
       return &regionBuilder;
     }
 
-    mlir::ArrayAttr iterator_types() {
+    mlir::ArrayAttr getIteratorTypes() {
       return getOperation()->getAttrOfType<mlir::ArrayAttr>("iterator_types");
     }
 
     mlir::ArrayAttr getIndexingMaps() {
       return getOperation()->getAttrOfType<mlir::ArrayAttr>("indexing_maps");
     }
 
     std::string getLibraryCallName() {
let extraClassDeclaration = [{
     }
 
     static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
                               mlir::ArrayRef<mlir::NamedAttribute>)>
     getRegionBuilder() {
       return &regionBuilder;
     }
 
-    mlir::ArrayAttr iterator_types() {
+    mlir::ArrayAttr getIteratorTypes() {
       return getOperation()->getAttrOfType<mlir::ArrayAttr>("iterator_types");
     }
 
     mlir::ArrayAttr getIndexingMaps() {
       return getOperation()->getAttrOfType<mlir::ArrayAttr>("indexing_maps");
     }
 
     std::string getLibraryCallName() {

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

value: !ScalarExpression
       scalar_fn:
         kind: type
         fn_name: cast_signed
         type_var: U
         operands:
         - !ScalarExpression
           scalar_arg: value
 
-#       IMPL:  Test3Op::iterator_types() {
+#       IMPL:  Test3Op::getIteratorTypesArray() {
 #  IMPL-NEXT:    int64_t rank = getRank(getOutputOperand(0));
 
 #       IMPL:  Test3Op::getIndexingMaps() {
 #  IMPL-NEXT:    MLIRContext *context = getContext();
 #  IMPL-NEXT:    AffineMap scalarMap = AffineMap::get(getNumParallelLoops(), 0, context);
 #  IMPL-NEXT:    AffineMap tensorMap = AffineMap::getMultiDimIdentityMap(
 
 # @linalg_structured_op

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

let builders = [
       {5}
     ];
     let hasCustomAssemblyFormat = 1;
     let hasFolder = 1;
     {6}
 
     let extraClassDeclaration = structuredOpsBaseDecls # [{{
       // Auto-generated.
-      ArrayAttr iterator_types();
+      SmallVector<StringRef> getIteratorTypesArray();
       ArrayAttr getIndexingMaps();
       static void regionBuilder(ImplicitLocOpBuilder &b,
                                 Block &block, ArrayRef<NamedAttribute> attrs);
       static std::function<void(ImplicitLocOpBuilder &,
                                 Block &, ArrayRef<NamedAttribute>)>
       getRegionBuilder() {{
         return regionBuilder;
       }
(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
        CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
   [{{
     {2}
     buildStructuredOp($_builder, $_state, resultTensorTypes, inputs, outputs,
       attributes, {0}::getRegionBuilder());
   }]>
 )FMT";
 
-// The iterator_types() method for structured ops. Parameters:
+// The getIteratorTypesArray() method for structured ops. Parameters:
 // {0}: Class name
 // {1}: Comma interleaved iterator type names.
 static const char structuredOpIteratorTypesFormat[] =
     R"FMT(
-ArrayAttr {0}::iterator_types() {{
-  return Builder(getContext()).getStrArrayAttr(SmallVector<StringRef>{{ {1} });
+SmallVector<StringRef> {0}::getIteratorTypesArray() {{
+  return SmallVector<StringRef>{{ {1} };
 }
 )FMT";
 
-// The iterator_types() method for rank polymorphic structured ops. Parameters:
+// The getIteratorTypesArray() method for rank polymorphic structured ops.
+// Parameters:
 // {0}: Class name
 static const char rankPolyStructuredOpIteratorTypesFormat[] =
     R"FMT(
-ArrayAttr {0}::iterator_types() {{
+SmallVector<StringRef> {0}::getIteratorTypesArray() {{
   int64_t rank = getRank(getOutputOperand(0));
-  return Builder(getContext()).getStrArrayAttr(
-    SmallVector<StringRef>(rank, getParallelIteratorTypeName()));
+  return SmallVector<StringRef>(rank, getParallelIteratorTypeName());
 }
 )FMT";
 
 // The indexing_maps() method for structured ops. Parameters:
 // {0}: Class name
 // {1}: Comma-separated list of dimension variable names.
 // {2}: Statements
 static const char structuredOpIndexingMapsFormat[] = R"FMT(