[mlir] Change IteratorType in ContractionOp in Vector dialect from string to enum.

[olegshyshkov]

Summary

This is the first step in replacing interator_type from strings with enums in Vector and Linalg dialect. This change adds IteratorTypeAttr and uses it in ContractionOp.

To avoid breaking all the tests, print/parse code has conversion between string and enum for now.

There is a shared code in StructuredOpsUtils.h that expects iterator types to be strings. To break this dependancy, this change forks helper function isParallelIterator and isReductionIterator to utils in both dialects and adds getIteratorTypeNames() to support backward compatibility with StructuredGenerator.

In the later changes, I plan to add a similar enum attribute to Linalg.

[pifon2a]

Very nice!

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

1491–1493

nit:

vector::IteratorType par = vector::IteratorType::parallel,
vector::IteratorType red = vector::IteratorType::reduction;

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

529–531

nit: omit {}

[olegshyshkov]

Fix nits.

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

let extraClassDeclaration = [{
     /// 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 llvm::to_vector(getIteratorTypes().getAsValueRange<StringAttr>());
+    }
+
     //========================================================================//
     // Forwarding functions to access interface methods from the
     // DestinationStyleOpInterface.
     // MLIR currently does not support dependent interfaces or interface
     // inheritance. By construction all ops with StructuredOpInterface must
     // implement DestinationStyleOpInterface.
     // TODO: reevalute the need for a cast when a better mechanism exists.
     //========================================================================//

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 
 /// Check if a LinalgOp is an element-wise operation.
 bool isElementwise(LinalgOp op);
 
 /// Check if `permutation` is a permutation of the range
 /// `[0, permutation.size())`.
 bool isPermutation(ArrayRef<int64_t> permutation);
 
+/// Check if `attr` has "parallel" iterator type semantics.
+bool isParallelIterator(Attribute attr);
+
+/// Check if `attr` has "reduction" iterator type semantics.
+bool isReductionIterator(Attribute attr);
+
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim);
 
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.

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

 /// 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"; }
 
 /// Use to encode that a particular iterator type has parallel semantics.
 constexpr StringRef getParallelIteratorTypeName() { return "parallel"; }
-inline bool isParallelIterator(Attribute attr) {
-  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
-  return strAttr && strAttr.getValue() == getParallelIteratorTypeName();
-}
 
 /// Use to encode that a particular iterator type has reduction semantics.
 constexpr StringRef getReductionIteratorTypeName() { return "reduction"; }
-inline bool isReductionIterator(Attribute attr) {
-  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
-  return strAttr && strAttr.getValue() == getReductionIteratorTypeName();
-}
 
 /// Use to encode that a particular iterator type has window semantics.
 constexpr StringRef getWindowIteratorTypeName() { return "window"; }
-inline bool isWindowIterator(Attribute attr) {
-  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
-  return strAttr && strAttr.getValue() == getWindowIteratorTypeName();
-}
 
 /// Use to encode that a particular iterator type has window semantics.
 inline ArrayRef<StringRef> getAllIteratorTypeNames() {
   static constexpr StringRef names[3] = {getParallelIteratorTypeName(),
                                          getReductionIteratorTypeName(),
                                          getWindowIteratorTypeName()};
   return llvm::makeArrayRef(names);
 }

 
 inline unsigned getNumIterators(ArrayAttr iteratorTypes) {
   unsigned res = 0;
   for (auto n : getAllIteratorTypeNames())
     res += getNumIterators(n, iteratorTypes);
   return res;
 }
 
-/// Typed representation for loop type strings.
-enum class IteratorType { Parallel, Reduction };
-
-inline StringRef toString(IteratorType t) {
-  switch (t) {
-  case IteratorType::Parallel:
-    return getParallelIteratorTypeName();
-  case IteratorType::Reduction:
-    return getReductionIteratorTypeName();
-  }
-  llvm_unreachable("Unsupported IteratorType");
-}
-
 /// Helper StructuredGenerator class to manipulate and rewrite ops with
 /// `StructuredOpInterface`. This is templated for now because VectorOps do not
 /// yet implement the StructuredOpInterface itself.
 template <typename StructuredOpInterface>
 class StructuredGenerator {
 public:
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
 
   struct IteratorType {
     IteratorType(StringRef strRef) : strRef(strRef) {}
-    bool isOfType(Attribute attr) const {
-      auto sAttr = attr.dyn_cast<StringAttr>();
-      return sAttr && sAttr.getValue() == strRef;
-    }
+    bool isOfType(StringRef typeName) const { return typeName == strRef; }
     StringRef strRef;
   };
   struct Par : public IteratorType {
     Par() : IteratorType(getParallelIteratorTypeName()) {}
   };
   struct Red : public IteratorType {
     Red() : IteratorType(getReductionIteratorTypeName()) {}
   };
   struct Win : public IteratorType {
     Win() : IteratorType(getWindowIteratorTypeName()) {}
   };
 
   StructuredGenerator(OpBuilder &builder, StructuredOpInterface op)
       : builder(builder), ctx(op.getContext()), loc(op.getLoc()),
-        iterators(op.getIteratorTypes()), maps(op.getIndexingMapsArray()),
+        iterators(op.getIteratorTypeNames()), maps(op.getIndexingMapsArray()),
         op(op) {}
 
   bool iters(ArrayRef<IteratorType> its) {
     if (its.size() != iterators.size())
       return false;
     for (int i = 0, e = its.size(); i != e; ++i) {
       if (!its[i].isOfType(iterators[i]))
         return false;
     }
     return true;
   }
 
   bool layout(MapList l) {
     auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
     return maps == infer(l);
   }
 
 protected:
   OpBuilder &builder;
   MLIRContext *ctx;
   Location loc;
-  ArrayAttr iterators;
+  SmallVector<StringRef> iterators;
   SmallVector<AffineMap, 4> maps;
   Operation *op;
 };
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_STRUCTUREDOPSUTILS_H

mlir/include/mlir/Dialect/Vector/IR/VectorOps.h

 /// memory.
 bool isDisjointTransferSet(VectorTransferOpInterface transferA,
                            VectorTransferOpInterface transferB);
 
 /// Return the result value of reducing two scalar/vector values with the
 /// corresponding arith operation.
 Value makeArithReduction(OpBuilder &b, Location loc, CombiningKind kind,
                          Value v1, Value v2);
+
+/// Returns true if `attr` has "parallel" iterator type semantics.
+inline bool isParallelIterator(Attribute attr) {
+  return attr.cast<IteratorTypeAttr>().getValue() == IteratorType::parallel;
+}
+
+/// Returns true if `attr` has "reduction" iterator type semantics.
+inline bool isReductionIterator(Attribute attr) {
+  return attr.cast<IteratorTypeAttr>().getValue() == IteratorType::reduction;
+}
+
 } // namespace vector
 } // namespace mlir
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Vector/IR/VectorOps.h.inc"
 #include "mlir/Dialect/Vector/IR/VectorOpsDialect.h.inc"
 
 #endif // MLIR_DIALECT_VECTOR_IR_VECTOROPS_H

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

 }
 
 /// An attribute that specifies the combining function for `vector.contract`,
 /// and `vector.reduction`.
 def Vector_CombiningKindAttr : EnumAttr<Vector_Dialect, CombiningKind, "kind"> {
   let assemblyFormat = "`<` $value `>`";
 }
 
+def IteratorType : I32EnumAttr<"IteratorType", "Iterator type", [
+  I32EnumAttrCase<"parallel", 0>,
+  I32EnumAttrCase<"reduction", 1>
+  ]> {
+    let genSpecializedAttr = 0;
+    let cppNamespace = "::mlir::vector";
+}
+
+def IteratorTypeEnum : EnumAttr<Vector_Dialect, IteratorType, "iterator_type"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
+def IteratorTypeArrayAttr : TypedArrayAttrBase<IteratorTypeEnum,
+    "Iterator type should be an enum.">;
+
 // TODO: Add an attribute to specify a different algebra with operators other
 // than the current set: {*, +}.
 def Vector_ContractionOp :
   Vector_Op<"contract", [
       NoSideEffect,
       PredOpTrait<"lhs and rhs have same element type", TCopVTEtIsSameAs<0, 1>>,
       PredOpTrait<"third operand acc and result have same element type",
                   TCresVTEtIsSameAsOpBase<0, 2>>,
       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
     ]>,
     Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyType:$acc,
                Variadic<VectorOf<[I1]>>:$masks,
                ArrayAttr:$indexing_maps,
-               ArrayAttr:$iterator_types,
+               IteratorTypeArrayAttr:$iterator_types,
                DefaultValuedAttr<Vector_CombiningKindAttr,
                                  "CombiningKind::ADD">:$kind)>,
     Results<(outs AnyType)> {
   let summary = "vector contraction operation";
   let description = [{
     Computes the sum of products of vector elements along contracting
     dimension pairs from 2 vectors of rank M and N respectively, adds this
     intermediate result to the accumulator argument of rank K, and returns a
%7 = vector.contract #contraction_trait %0, %1, %2
       : vector<10xf32>, vector<10xf32> into f32
     ```
   }];
   let builders = [
     OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
       "ArrayAttr":$indexingMaps, "ArrayAttr":$iteratorTypes)>,
     OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
       "ArrayRef<ArrayRef<AffineExpr>>":$indexingExprs,
-      "ArrayRef<StringRef>":$iteratorTypes)>,
+      "ArrayRef<IteratorType>":$iteratorTypes)>,
     OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
       "ArrayAttr":$indexingMaps, "ArrayAttr":$iteratorTypes,
       "CombiningKind":$kind)>
   ];
   let extraClassDeclaration = [{
     VectorType getLhsType() {
       return getLhs().getType().cast<VectorType>();
     }
let extraClassDeclaration = [{
     std::vector<std::pair<int64_t, int64_t>> getContractingDimMap();
     std::vector<std::pair<int64_t, int64_t>> getBatchDimMap();
 
     static constexpr StringRef getKindAttrStrName() { return "kind"; }
 
     static CombiningKind getDefaultKind() {
       return CombiningKind::ADD;
     }
+
+    // Returns iterator types in string format.
+    SmallVector<StringRef> getIteratorTypeNames() {
+      return llvm::to_vector(
+          llvm::map_range(getIteratorTypes(), [](Attribute a) {
+            return stringifyIteratorType(a.cast<IteratorTypeAttr>().getValue());
+          }));
+    }
   }];
 
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def Vector_ReductionOp :

mlir/lib/Conversion/VectorToGPU/NvGpuSupport.h

 FailureOr<AffineMap>
 getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
                                   const WarpMatrixInfo &fragmentType);
 
 struct LdMatrixParams {
   VectorType fragmentType;
   bool isAccum;
   int64_t numTiles;
-  IteratorType contiguousDimType;
+  vector::IteratorType contiguousDimType;
   NVVM::MMALayout targetLayout;
 };
 
 FailureOr<LdMatrixParams> getLdMatrixParams(const WarpMatrixInfo &type,
                                             bool transpose);
 /// Returns an AffineMap which maps a single dimension representing the laneId
 /// to two results representing offsets within the matrix operand that should
 /// be the pointer locations a thread should pass to the ldmatrix instruction.

mlir/lib/Conversion/VectorToGPU/NvGpuSupport.cpp

FailureOr<nvgpu::LdMatrixParams> getLdMatrixParams(const WarpMatrixInfo &type,
   params.fragmentType = type.vectorType;
   if (type.operandRole == MatMulOperandRole::A ||
       type.operandRole == MatMulOperandRole::C) {
     params.targetLayout = NVVM::MMALayout::row;
   } else {
     params.targetLayout = NVVM::MMALayout::col;
   }
   ArrayRef<int64_t> shape = type.vectorType.getShape();
-  params.contiguousDimType =
-      transpose ? IteratorType::Parallel : IteratorType::Reduction;
+  params.contiguousDimType = transpose ? vector::IteratorType::parallel
+                                       : vector::IteratorType::reduction;
 
-  if (params.contiguousDimType == IteratorType::Reduction) {
+  if (params.contiguousDimType == vector::IteratorType::reduction) {
     params.numTiles = (shape[0] / kNumRowsPerTile) *
                       ((shape[1] * elType.getIntOrFloatBitWidth()) / 128);
   } else {
     params.numTiles = (shape[1] / kNumRowsPerTile) *
                       ((shape[0] * elType.getIntOrFloatBitWidth()) / 128);
   }
 
   if (params.numTiles == 0)
getLaneIdToLdMatrixMatrixCoord(Location loc, OpBuilder &builder,
   ArrayRef<int64_t> operandShape = params.fragmentType.getShape();
   AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
 
   auto makeMap = [&](ArrayRef<AffineExpr> dimExprs) -> AffineMap {
     return AffineMap::get(1, 0, dimExprs, builder.getContext());
   };
 
   // This case corresponds to row-major A|C or col-major B operands.
-  if (params.contiguousDimType == IteratorType::Reduction) {
+  if (params.contiguousDimType == vector::IteratorType::reduction) {
     AffineExpr row = d0 % (operandShape[0]);
     AffineExpr col = d0.floorDiv(operandShape[0]) * (kElementsPer128b);
     return makeMap({row, col});
   }
 
   // This case Corresponds to col-major A|C or row-major B operands. The
   // operandShape given is already pre-transposed (e.g. 8x16 = KxN).
-  if (params.contiguousDimType == IteratorType::Parallel) {
+  if (params.contiguousDimType == vector::IteratorType::parallel) {
     const int64_t num8x128bCols = (operandShape[0] * bitsPerElement) / 128;
     // Threads are assigned in groups of 8 first across columns, then to
     // rows. This is transpose of what `ldmatrix` expects, but when
     // `ldmatrix` gets the `.trans` qualifier, final the effect will be to
     // transpose just the blocks.
     auto groupIdx = d0.floorDiv(kNumThreadsPerTile);
     auto tileCol = (groupIdx % num8x128bCols);
     auto tileRow = groupIdx.floorDiv(num8x128bCols);
PrepareContractToGPUMMASync::matchAndRewrite(vector::ContractionOp op,
   AffineExpr n;
   AffineExpr k;
   bindDims(rewriter.getContext(), m, n, k);
   static constexpr std::array<int64_t, 2> perm = {1, 0};
   auto iteratorTypes = op.getIteratorTypes().getValue();
   SmallVector<AffineMap, 4> maps = op.getIndexingMapsArray();
   if (iteratorTypes.size() != 3)
     return failure();
-  if (!(isParallelIterator(iteratorTypes[0]) &&
-        isParallelIterator(iteratorTypes[1]) &&
-        isReductionIterator(iteratorTypes[2])))
+  if (!(vector::isParallelIterator(iteratorTypes[0]) &&
+        vector::isParallelIterator(iteratorTypes[1]) &&
+        vector::isReductionIterator(iteratorTypes[2])))
     return failure();
 
   // The canonical form is "TNT" = A row-major, B col-major, C row-major.
   const auto canonicalForm = infer({{m, k}, {n, k}, {m, n}});
   if (maps == canonicalForm) {
     return failure();
   }
   if (maps == infer({{m, k}, {k, n}, {m, n}})) {

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

static bool contractSupportsMMAMatrixType(vector::ContractionOp contract,
   if (llvm::size(contract.getMasks()) != 0)
     return false;
 
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
   AffineExpr m, n, k;
   bindDims(contract.getContext(), m, n, k);
   auto iteratorTypes = contract.getIteratorTypes().getValue();
-  if (!(isParallelIterator(iteratorTypes[0]) &&
-        isParallelIterator(iteratorTypes[1]) &&
-        isReductionIterator(iteratorTypes[2])))
+  if (!(vector::isParallelIterator(iteratorTypes[0]) &&
+        vector::isParallelIterator(iteratorTypes[1]) &&
+        vector::isReductionIterator(iteratorTypes[2])))
     return false;
 
   // The contract needs to represent a matmul to be able to convert to
   // MMAMatrix matmul.
   if (!useNvGpu &&
       contract.getIndexingMapsArray() != infer({{m, k}, {k, n}, {m, n}}))
     return false;
   if (useNvGpu &&
LogicalResult matchAndRewrite(vector::ContractionOp op,
     // Set up the parallel/reduction structure in right form.
     using MapList = ArrayRef<ArrayRef<AffineExpr>>;
     auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
     AffineExpr m, n, k;
     bindDims(rewriter.getContext(), m, n, k);
     static constexpr std::array<int64_t, 2> perm = {1, 0};
     auto iteratorTypes = op.getIteratorTypes().getValue();
     SmallVector<AffineMap, 4> maps = op.getIndexingMapsArray();
-    if (!(isParallelIterator(iteratorTypes[0]) &&
-          isParallelIterator(iteratorTypes[1]) &&
-          isReductionIterator(iteratorTypes[2])))
+    if (!(vector::isParallelIterator(iteratorTypes[0]) &&
+          vector::isParallelIterator(iteratorTypes[1]) &&
+          vector::isReductionIterator(iteratorTypes[2])))
       return failure();
     //
     // Two outer parallel, one inner reduction (matmat flavor).
     //
     if (maps == infer({{m, k}, {k, n}, {m, n}})) {
       // This is the classical row-major matmul, nothing to do.
       return failure();
     }

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

return builder
         .create<vector::TransferWriteOp>(loc, res, resShaped,
                                          ValueRange{zero, zero, zero})
         .getOperation();
   }
 
   // Create a contraction: lhs{n, w, c} * rhs{c, f} -> res{n, w, f}
   Value conv1dSliceAsContraction(OpBuilder &b, Location loc, Value lhs,
                                  Value rhs, Value res) {
-    StringRef par = Par().strRef, red = Red().strRef;
+    vector::IteratorType par = vector::IteratorType::parallel;
+    vector::IteratorType red = vector::IteratorType::reduction;
     AffineExpr n, w, f, c;
     bindDims(ctx, n, w, f, c);
     return builder.create<vector::ContractionOp>(
         loc, lhs, rhs, res,
         /*indexingMaps=*/MapList{{n, w, c}, {c, f}, {n, w, f}},
-        /*iteratorTypes=*/ArrayRef<StringRef>{par, par, par, red});
+        /*iteratorTypes=*/ArrayRef<vector::IteratorType>{par, par, par, red});
   }
 
   /// Generate a vector implementation for:
   /// ```
   ///   Op def: (     n,     w,     c,    kw)
   ///    Iters: ({Par(), Par(), Par(), Red()})
   ///   Layout: {{n, strideW * w + dilationW * kw, c}, {kw, c}, {n, w, c}}
   /// ```

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

for (auto index : permutation) {
     if (index < 0 || index >= static_cast<int64_t>(permutation.size()))
       return false;
     indexCounts[index]++;
   }
   // Return true if all indices appear once.
   return count(indexCounts, 1) == static_cast<int64_t>(permutation.size());
 }
 
+bool isParallelIterator(Attribute attr) {
+  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
+  return strAttr && strAttr.getValue() == getParallelIteratorTypeName();
+}
+
+bool isReductionIterator(Attribute attr) {
+  auto strAttr = attr.dyn_cast_or_null<StringAttr>();
+  return strAttr && strAttr.getValue() == getReductionIteratorTypeName();
+}
+
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");

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

static bool isAdmissableTensorExp(Merger &merger, linalg::GenericOp op,
   // [Bik96,Ch9], is also admissable without special codegen.
   if (merger.isSingleCondition(tensor, exp))
     return true;
   // Accept "truly dynamic" if the output tensor materializes uninitialized
   // into the computation and insertions occur in lexicographic index order.
   if (isMaterializing(lhs->get())) {
     unsigned nest = 0;
     for (unsigned i = 0; i < numLoops; i++) {
-      if (isReductionIterator(iteratorTypes[topSort[i]]))
+      if (linalg::isReductionIterator(iteratorTypes[topSort[i]]))
         break; // terminate at first reduction
       nest++;
     }
     // Determine admissable dynamic insertion situations:
     // (1) fully injective, since there are no reductions,
     // (2) admissable 1-d expansion in innermost dimension.
     if (nest >= op.getRank(lhs) - 1) {
       *sparseOut = lhs;

 /// Generates a for-loop on a single index.
 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
                          unsigned idx, BitVector &indices) {
   unsigned fb = indices.find_first();
   unsigned tensor = merger.tensor(fb);
   assert(idx == merger.index(fb));
   auto iteratorTypes = op.iterator_types().getValue();
-  bool isReduction = isReductionIterator(iteratorTypes[idx]);
+  bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
   bool isSparse = merger.isDim(fb, Dim::kSparse);
   bool isVector = isVectorFor(codegen, isInner, isReduction, isSparse) &&
                   denseUnitStrides(merger, op, idx);
   bool isParallel =
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
   // Prepare vector length.
   if (isVector)

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 
 //===----------------------------------------------------------------------===//
 // ContractionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
                                   ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
-                                  ArrayRef<StringRef> iteratorTypes) {
+                                  ArrayRef<IteratorType> iteratorTypes) {
   result.addOperands({lhs, rhs, acc});
   result.addTypes(acc.getType());
   result.addAttribute(::mlir::getIndexingMapsAttrName(),
                       builder.getAffineMapArrayAttr(
                           AffineMap::inferFromExprList(indexingExprs)));
-  result.addAttribute(::mlir::getIteratorTypesAttrName(),
-                      builder.getStrArrayAttr(iteratorTypes));
+  result.addAttribute(
+      ::mlir::getIteratorTypesAttrName(),
+      builder.getArrayAttr(llvm::to_vector(llvm::map_range(
+          iteratorTypes, [&](IteratorType t) -> mlir::Attribute {
+            return IteratorTypeAttr::get(builder.getContext(), t);
+          }))));
 }
 
 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
                                   ArrayAttr indexingMaps,
                                   ArrayAttr iteratorTypes) {
   build(builder, result, lhs, rhs, acc, indexingMaps, iteratorTypes,
         ContractionOp::getDefaultKind());
if (parser.parseAttribute(dictAttr, "_", result.attributes) ||
       parser.parseKeywordType("into", resultType) ||
       parser.resolveOperand(lhsInfo, types[0], result.operands) ||
       parser.resolveOperand(rhsInfo, types[1], result.operands) ||
       parser.resolveOperand(accInfo, resultType, result.operands) ||
       parser.addTypeToList(resultType, result.types))
     return failure();
   result.attributes.assign(dictAttr.getValue().begin(),
                            dictAttr.getValue().end());
+
+  // Convert array of string into an array of IteratyType enums. This is needed,
+  // because tests still use the old format when 'iterator_types' attribute is
+  // represented as an array of strings.
+  // TODO: Remove this conversion once tests are fixed.
+  ArrayAttr iteratorTypes =
+      result.attributes.get("iterator_types").cast<ArrayAttr>();
+
+  SmallVector<Attribute> iteratorTypeAttrs;
+
+  for (StringRef s : iteratorTypes.getAsValueRange<StringAttr>()) {
+    auto maybeIteratorType = symbolizeIteratorType(s);
+    if (!maybeIteratorType.hasValue())
+      return parser.emitError(loc) << "unexpected iterator_type (" << s << ")";
+
+    iteratorTypeAttrs.push_back(IteratorTypeAttr::get(
+        parser.getContext(), maybeIteratorType.getValue()));
+  }
+  result.attributes.set("iterator_types",
+                        parser.getBuilder().getArrayAttr(iteratorTypeAttrs));
+
   if (!result.attributes.get(ContractionOp::getKindAttrStrName())) {
     result.addAttribute(
         ContractionOp::getKindAttrStrName(),
         CombiningKindAttr::get(result.getContext(),
                                ContractionOp::getDefaultKind()));
   }
   if (masksInfo.empty())
     return success();

 }
 
 void ContractionOp::print(OpAsmPrinter &p) {
   // TODO: Unify printing code with linalg ops.
   auto attrNames = getTraitAttrNames();
   llvm::StringSet<> traitAttrsSet;
   traitAttrsSet.insert(attrNames.begin(), attrNames.end());
   SmallVector<NamedAttribute, 8> attrs;
-  for (auto attr : (*this)->getAttrs())
-    if (traitAttrsSet.count(attr.getName().strref()) > 0)
+  for (auto attr : (*this)->getAttrs()) {
+    if (attr.getName() == getIteratorTypesAttrName()) {
+      auto iteratorTypes =
+          attr.getValue()
+              .cast<ArrayAttr>()
+              .getAsValueRange<IteratorTypeAttr, IteratorType>();
+      // Convert IteratorType enums into the string representation. This is
+      // needed, because tests still use the old format when 'iterator_types'
+      // attribute is represented as an array of strings.
+      // TODO: Remove this conversion once tests are fixed.
+      SmallVector<Attribute> iteratorTypeNames = llvm::to_vector(
+          llvm::map_range(iteratorTypes, [&](IteratorType t) -> Attribute {
+            return StringAttr::get(getContext(), stringifyIteratorType(t));
+          }));
+
+      attrs.emplace_back(getIteratorTypesAttrName(),
+                         ArrayAttr::get(getContext(), iteratorTypeNames));
+    } else if (traitAttrsSet.count(attr.getName().strref()) > 0)
       attrs.push_back(attr);
+  }
 
   auto dictAttr = DictionaryAttr::get(getContext(), attrs);
   p << " " << dictAttr << " " << getLhs() << ", ";
   p << getRhs() << ", " << getAcc();
   if (getMasks().size() == 2)
     p << ", " << getMasks();
 
   p.printOptionalAttrDict((*this)->getAttrs(), attrNames);
static int64_t getResultIndex(AffineMap map, AffineExpr targetExpr) {
   for (int64_t i = 0, e = map.getNumResults(); i < e; ++i)
     if (targetExpr == map.getResult(i))
       return i;
   return -1;
 }
 
 static std::vector<std::pair<int64_t, int64_t>>
 getDimMap(ArrayRef<AffineMap> indexingMaps, ArrayAttr iteratorTypes,
-          StringRef targetIteratorTypeName, MLIRContext *context) {
+          IteratorType targetIteratorType, MLIRContext *context) {
   std::vector<std::pair<int64_t, int64_t>> dimMap;
   for (const auto &it : llvm::enumerate(iteratorTypes)) {
-    auto iteratorTypeName = it.value().cast<StringAttr>().getValue();
-    if (iteratorTypeName != targetIteratorTypeName)
+    auto iteratorType = it.value().cast<IteratorTypeAttr>().getValue();
+    if (iteratorType != targetIteratorType)
       continue;
     // Search lhs/rhs map results for 'targetExpr'.
     auto targetExpr = getAffineDimExpr(it.index(), context);
     int64_t lhsDim = getResultIndex(indexingMaps[0], targetExpr);
     int64_t rhsDim = getResultIndex(indexingMaps[1], targetExpr);
     if (lhsDim >= 0 && rhsDim >= 0)
       dimMap.emplace_back(lhsDim, rhsDim);
   }
   return dimMap;
 }
 
 void ContractionOp::getIterationBounds(
     SmallVectorImpl<int64_t> &iterationBounds) {
   auto lhsShape = getLhsType().getShape();
   auto resVectorType = getResultType().dyn_cast<VectorType>();
   SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
   SmallVector<int64_t, 2> iterationShape;
   for (const auto &it : llvm::enumerate(getIteratorTypes())) {
     // Search lhs/rhs map results for 'targetExpr'.
     auto targetExpr = getAffineDimExpr(it.index(), getContext());
-    auto iteratorTypeName = it.value().cast<StringAttr>().getValue();
-    if (iteratorTypeName == getReductionIteratorTypeName()) {
+    auto iteratorType = it.value().cast<IteratorTypeAttr>().getValue();
+    if (iteratorType == IteratorType::reduction) {
       // Get reduction dim size from lhs shape (same size in rhsShape).
       int64_t lhsDimIndex = getResultIndex(indexingMaps[0], targetExpr);
       assert(lhsDimIndex >= 0);
       iterationBounds.push_back(lhsShape[lhsDimIndex]);
       continue;
     }
     // Get parallel dimension size from result shape.
     int64_t resDimIndex = getResultIndex(indexingMaps[2], targetExpr);
for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
       auto dim = map.getResult(i).cast<AffineDimExpr>();
       iterationIndexMap[index][dim.getPosition()] = i;
     }
   }
 }
 
 std::vector<std::pair<int64_t, int64_t>> ContractionOp::getContractingDimMap() {
   SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
-  return getDimMap(indexingMaps, getIteratorTypes(),
-                   getReductionIteratorTypeName(), getContext());
+  return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::reduction,
+                   getContext());
 }
 
 std::vector<std::pair<int64_t, int64_t>> ContractionOp::getBatchDimMap() {
   SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
-  return getDimMap(indexingMaps, getIteratorTypes(),
-                   getParallelIteratorTypeName(), getContext());
+  return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::parallel,
+                   getContext());
 }
 
 Optional<SmallVector<int64_t, 4>> ContractionOp::getShapeForUnroll() {
   SmallVector<int64_t, 4> shape;
   getIterationBounds(shape);
   return shape;
 }
 

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 #include "llvm/ADT/DenseSet.h"
LogicalResult matchAndRewrite(vector::MultiDimReductionOp reduceOp,
     if (reduceOp.getKind() != vector::CombiningKind::ADD)
       return failure();
     Operation *mulOp = reduceOp.getSource().getDefiningOp();
     if (!mulOp || !isa<arith::MulIOp, arith::MulFOp>(mulOp))
       return failure();
     SmallVector<bool> reductionMask = reduceOp.getReductionMask();
     auto srcMap = rewriter.getMultiDimIdentityMap(reductionMask.size());
     SmallVector<AffineExpr> exprs;
-    SmallVector<StringRef> iteratorTypes;
+    SmallVector<vector::IteratorType> iteratorTypes;
     for (const auto &isReduceDim : llvm::enumerate(reductionMask)) {
       if (!isReduceDim.value()) {
-        iteratorTypes.push_back(getParallelIteratorTypeName());
+        iteratorTypes.push_back(vector::IteratorType::parallel);
         exprs.push_back(rewriter.getAffineDimExpr(isReduceDim.index()));
       } else {
-        iteratorTypes.push_back(getReductionIteratorTypeName());
+        iteratorTypes.push_back(vector::IteratorType::reduction);
       }
     }
     auto dstMap = AffineMap::get(/*dimCount=*/reductionMask.size(),
                                  /*symCount=*/0, exprs, reduceOp.getContext());
     rewriter.replaceOpWithNewOp<mlir::vector::ContractionOp>(
         reduceOp, mulOp->getOperand(0), mulOp->getOperand(1), reduceOp.getAcc(),
         rewriter.getAffineMapArrayAttr({srcMap, srcMap, dstMap}),
-        rewriter.getStrArrayAttr(iteratorTypes));
+        rewriter.getArrayAttr(llvm::to_vector(llvm::map_range(
+            iteratorTypes, [&](IteratorType t) -> mlir::Attribute {
+              return IteratorTypeAttr::get(rewriter.getContext(), t);
+            }))));
     return success();
   }
 };
 
 /// Merge TransposeOp into ContractionOp user.
 /// Ex:
 /// ```
 ///   %0 = vector.transpose %arg0, [2, 0, 1]