[mlir] NFC - Refactor and expose a parsing helper for OffsetSizeAndStrideInterface

[nicolasvasilache]

Summary

Parse trailing part of an op of the form:

<optional-offset-prefix>`[` offset-list `]`
<optional-size-prefix>`[` size-list `]`
<optional-stride-prefix>[` stride-list `]`

Each entry in the offset, size and stride list either resolves to an integer
constant or an operand of index type.
Constants are added to the result as named integer array attributes with
name OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName() (resp.
getStaticSizesAttrName(), getStaticStridesAttrName()).

Append the number of offset, size and stride operands to segmentSizes
before adding it to result as the named attribute:
OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr().
Offset, size and stride operands resolution occurs after preResolutionFn
to give a chance to leading operands to resolve first, after parsing the
types.

ParseResult parseOffsetsSizesAndStrides(
    OpAsmParser &parser, OperationState &result, ArrayRef<int> segmentSizes,
    llvm::function_ref<ParseResult(OpAsmParser &, OperationState &)>
        preResolutionFn = nullptr,
    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalOffsetPrefix =
        nullptr,
    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalSizePrefix =
        nullptr,
    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalStridePrefix =
        nullptr);

[pifon2a]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3175

here as well

3774

nit: inline in parseOffsetsSizesAndStrides call?

3882

and here

mlir/lib/Interfaces/ViewLikeInterface.cpp

138

nit: move auto b = ... above
auto indexType = b.getIndexType().

[nicolasvasilache]

Thanks pifon@, made your suggestions and squashed in the followup commit.

mlir/lib/Interfaces/ViewLikeInterface.cpp

138

Actually, just inlined it.

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 //===- ViewLikeInterface.h - View-like operations interface ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the operation interface for view-like operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
 #define MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
 
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
 
 namespace mlir {
 /// Auxiliary range data structure to unpack the offset, size and stride
 /// operands into a list of triples. Such a list can be more convenient to
 /// manipulate.
 struct Range {
   Value offset;
   Value size;
   Value stride;
 };
 
 class OffsetSizeAndStrideOpInterface;
 LogicalResult verify(OffsetSizeAndStrideOpInterface op);
+
+/// Parse trailing part of an op of the form:
+/// ```
+///   <optional-offset-prefix>`[` offset-list `]`
+///   <optional-size-prefix>`[` size-list `]`
+///   <optional-stride-prefix>[` stride-list `]`
+/// ```
+/// Each entry in the offset, size and stride list either resolves to an integer
+/// constant or an operand of index type.
+/// Constants are added to the `result` as named integer array attributes with
+/// name `OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName()` (resp.
+/// `getStaticSizesAttrName()`, `getStaticStridesAttrName()`).
+///
+/// Append the number of offset, size and stride operands to `segmentSizes`
+/// before adding it to `result` as the named attribute:
+/// `OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr()`.
+///
+/// Offset, size and stride operands resolution occurs after `preResolutionFn`
+/// to give a chance to leading operands to resolve first, after parsing the
+/// types.
+ParseResult parseOffsetsSizesAndStrides(
+    OpAsmParser &parser, OperationState &result, ArrayRef<int> segmentSizes,
+    llvm::function_ref<ParseResult(OpAsmParser &, OperationState &)>
+        preResolutionFn = nullptr,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalOffsetPrefix =
+        nullptr,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalSizePrefix =
+        nullptr,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalStridePrefix =
+        nullptr);
+
 } // namespace mlir
 
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/ViewLikeInterface.h.inc"
 
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE_H_

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

llvm::interleaveComma(arrayAttr, p, [&](Attribute a) {
     if (isDynamic(val))
       p << values[idx++];
     else
       p << val;
   });
   p << ']';
 }
 
-/// Parse a mixed list with either (1) static integer values or (2) SSA values.
-/// Fill `result` with the integer ArrayAttr named `attrName` where `dynVal`
-/// encode the position of SSA values. Add the parsed SSA values to `ssa`
-/// in-order.
-//
-/// E.g. after parsing "[%arg0, 7, 42, %arg42]":
-///   1. `result` is filled with the i64 ArrayAttr "[`dynVal`, 7, 42, `dynVal`]"
-///   2. `ssa` is filled with "[%arg0, %arg1]".
-static ParseResult
-parseListOfOperandsOrIntegers(OpAsmParser &parser, OperationState &result,
-                              StringRef attrName, int64_t dynVal,
-                              SmallVectorImpl<OpAsmParser::OperandType> &ssa) {
-  if (failed(parser.parseLSquare()))
-    return failure();
-  // 0-D.
-  if (succeeded(parser.parseOptionalRSquare())) {
-    result.addAttribute(attrName, parser.getBuilder().getArrayAttr({}));
-    return success();
-  }
-
-  SmallVector<int64_t, 4> attrVals;
-  while (true) {
-    OpAsmParser::OperandType operand;
-    auto res = parser.parseOptionalOperand(operand);
-    if (res.hasValue() && succeeded(res.getValue())) {
-      ssa.push_back(operand);
-      attrVals.push_back(dynVal);
-    } else {
-      IntegerAttr attr;
-      if (failed(parser.parseAttribute<IntegerAttr>(attr)))
-        return parser.emitError(parser.getNameLoc())
-               << "expected SSA value or integer";
-      attrVals.push_back(attr.getInt());
-    }
-
-    if (succeeded(parser.parseOptionalComma()))
-      continue;
-    if (failed(parser.parseRSquare()))
-      return failure();
-    break;
-  }
-
-  auto arrayAttr = parser.getBuilder().getI64ArrayAttr(attrVals);
-  result.addAttribute(attrName, arrayAttr);
-  return success();
-}
-
 /// Verify that a particular offset/size/stride static attribute is well-formed.
 static LogicalResult verifyOpWithOffsetSizesAndStridesPart(
     OffsetSizeAndStrideOpInterface op, StringRef name,
     unsigned expectedNumElements, StringRef attrName, ArrayAttr attr,
     llvm::function_ref<bool(int64_t)> isDynamic, ValueRange values) {
   /// Check static and dynamic offsets/sizes/strides breakdown.
   if (attr.size() != expectedNumElements)
     return op.emitError("expected ")
void mlir::MemRefReinterpretCastOp::build(
     ValueRange strides, ArrayRef<NamedAttribute> attrs) {
   build(b, result, resultType, source, offset, sizes, strides,
         b.getI64ArrayAttr(staticOffset), b.getI64ArrayAttr(staticSizes),
         b.getI64ArrayAttr(staticStrides));
   result.addAttributes(attrs);
 }
 
 /// Build a MemRefReinterpretCastOp with all dynamic entries: `staticOffsets`,
-/// `staticSizes` and `staticStrides` are  automatically filled with
+/// `staticSizes` and `staticStrides` are automatically filled with
 /// source-memref-rank sentinel values that encode dynamic entries.
 void mlir::MemRefReinterpretCastOp::build(OpBuilder &b, OperationState &result,
                                           MemRefType resultType, Value source,
                                           Value offset, ValueRange sizes,
                                           ValueRange strides,
                                           ArrayRef<NamedAttribute> attrs) {
   unsigned rank = resultType.getRank();
   SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamicSize);
static void print(OpAsmPrinter &p, MemRefReinterpretCastOp op) {
   printListOfOperandsOrIntegers(p, op.strides(), op.static_strides(),
                                 ShapedType::isDynamicStrideOrOffset);
   p.printOptionalAttrDict(
       op.getAttrs(),
       /*elidedAttrs=*/MemRefReinterpretCastOp::getSpecialAttrNames());
   p << ": " << op.source().getType() << " to " << op.getType();
 }
 
-/// Parse of the form:
+/// Parse a memref_reinterpret_cast op of the form:
 /// ```
-///   `name` ssa-name to
+///   `memref_reinterpret_cast` ssa-name to
 ///       offset: `[` offset `]`
 ///       sizes: `[` size-list `]`
 ///       strides:`[` stride-list `]`
 ///   `:` any-memref-type to strided-memref-type
 /// ```
 static ParseResult parseMemRefReinterpretCastOp(OpAsmParser &parser,
                                                 OperationState &result) {
-  // Parse `operand` and `offset`.
-  OpAsmParser::OperandType operand;
-  if (parser.parseOperand(operand))
-    return failure();
-
-  // Parse offset.
-  SmallVector<OpAsmParser::OperandType, 1> offset;
-  if (parser.parseKeyword("to") || parser.parseKeyword("offset") ||
-      parser.parseColon() ||
-      parseListOfOperandsOrIntegers(
-          parser, result,
-          OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName(),
-          ShapedType::kDynamicStrideOrOffset, offset) ||
-      parser.parseComma())
-    return failure();
-
-  // Parse `sizes`.
-  SmallVector<OpAsmParser::OperandType, 4> sizes;
-  if (parser.parseKeyword("sizes") || parser.parseColon() ||
-      parseListOfOperandsOrIntegers(
-          parser, result,
-          OffsetSizeAndStrideOpInterface::getStaticSizesAttrName(),
-          ShapedType::kDynamicSize, sizes) ||
-      parser.parseComma())
-    return failure();
-
-  // Parse `strides`.
-  SmallVector<OpAsmParser::OperandType, 4> strides;
-  if (parser.parseKeyword("strides") || parser.parseColon() ||
-      parseListOfOperandsOrIntegers(
-          parser, result,
-          OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
-          ShapedType::kDynamicStrideOrOffset, strides))
-    return failure();
-
-  // Handle segment sizes.
-  auto b = parser.getBuilder();
-  SmallVector<int, 4> segmentSizes = {1, static_cast<int>(offset.size()),
-                                      static_cast<int>(sizes.size()),
-                                      static_cast<int>(strides.size())};
-  result.addAttribute(MemRefReinterpretCastOp::getOperandSegmentSizeAttr(),
-
-                      b.getI32VectorAttr(segmentSizes));
-
-  // Parse types and resolve.
-  Type indexType = b.getIndexType();
-  Type operandType, resultType;
-  return failure(
-      (parser.parseOptionalAttrDict(result.attributes) ||
-       parser.parseColonType(operandType) || parser.parseKeyword("to") ||
-       parser.parseType(resultType) ||
-       parser.resolveOperand(operand, operandType, result.operands) ||
-       parser.resolveOperands(offset, indexType, result.operands) ||
-       parser.resolveOperands(sizes, indexType, result.operands) ||
-       parser.resolveOperands(strides, indexType, result.operands) ||
-       parser.addTypeToList(resultType, result.types)));
+  // Parse `operand`
+  OpAsmParser::OperandType srcInfo;
+  if (parser.parseOperand(srcInfo))
+    return failure();
+
+  auto parseOffsetPrefix = [](OpAsmParser &parser) {
+    return failure(parser.parseKeyword("to") || parser.parseKeyword("offset") ||
+                   parser.parseColon());
+  };
+  auto parseSizePrefix = [](OpAsmParser &parser) {
+    return failure(parser.parseComma() || parser.parseKeyword("sizes") ||
+                   parser.parseColon());
+  };
+  auto parseStridePrefix = [](OpAsmParser &parser) {
+    return failure(parser.parseComma() || parser.parseKeyword("strides") ||
+                   parser.parseColon());
+  };
+
+  Type srcType, dstType;
+  auto preResolutionFn = [&](OpAsmParser &parser, OperationState &result) {
+    return failure(parser.parseOptionalAttrDict(result.attributes) ||
+                   parser.parseColonType(srcType) ||
+                   parser.parseKeywordType("to", dstType) ||
+                   parser.resolveOperand(srcInfo, srcType, result.operands));
+  };
+  SmallVector<int, 4> segmentSizes{1}; // source memref
+  if (failed(parseOffsetsSizesAndStrides(parser, result, segmentSizes,
+                                         preResolutionFn, parseOffsetPrefix,
+                                         parseSizePrefix, parseStridePrefix)))
+    return failure();
+  return parser.addTypeToList(dstType, result.types);
 }
 
 static LogicalResult verify(MemRefReinterpretCastOp op) {
   // The source and result memrefs should be in the same memory space.
   auto srcType = op.source().getType().cast<BaseMemRefType>();
   auto resultType = op.getType().cast<MemRefType>();
   if (srcType.getMemorySpace() != resultType.getMemorySpace())
     return op.emitError("different memory spaces specified for source type ")
static void printOpWithOffsetsSizesAndStrides(
   p << " : " << op.getSourceType() << " " << resultTypeKeyword << " "
     << op.getType();
 }
 
 static void print(OpAsmPrinter &p, SubViewOp op) {
   return printOpWithOffsetsSizesAndStrides<SubViewOp>(p, op);
 }
 
-/// Parse of the form:
+/// Parse a subview op of the form:
 /// ```
-///   `name` ssa-name (extra-operands)?
-///       `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
-///     `:` strided-memref-type `resultTypeKeyword strided-memref-type
+///   `subview` ssa-name
+///     `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
+///     `:` strided-memref-type `to` strided-memref-type
 /// ```
-template <typename OpType>
-static ParseResult parseOpWithOffsetsSizesAndStrides(
-    OpAsmParser &parser, OperationState &result,
-    std::function<ParseResult(OpAsmParser &p,
-                              OpAsmParser::OperandType &dstInfo)>
-        parseExtraOperand = nullptr,
-    StringRef resultTypeKeyword = "to") {
-  OpAsmParser::OperandType srcInfo, dstInfo;
-  SmallVector<OpAsmParser::OperandType, 4> offsetsInfo, sizesInfo, stridesInfo;
-  auto indexType = parser.getBuilder().getIndexType();
-  Type srcType, dstType;
+static ParseResult parseSubViewOp(OpAsmParser &parser, OperationState &result) {
+  OpAsmParser::OperandType srcInfo;
   if (parser.parseOperand(srcInfo))
     return failure();
-  if (parseExtraOperand && parseExtraOperand(parser, dstInfo))
-    return failure();
-  if (parseListOfOperandsOrIntegers(
-          parser, result,
-          OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName(),
-          ShapedType::kDynamicStrideOrOffset, offsetsInfo) ||
-      parseListOfOperandsOrIntegers(
-          parser, result,
-          OffsetSizeAndStrideOpInterface::getStaticSizesAttrName(),
-          ShapedType::kDynamicSize, sizesInfo) ||
-      parseListOfOperandsOrIntegers(
-          parser, result,
-          OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
-          ShapedType::kDynamicStrideOrOffset, stridesInfo))
-    return failure();
-
-  // Handle segment sizes.
-  auto b = parser.getBuilder();
-  SmallVector<int, 4> segmentSizes = {1, static_cast<int>(offsetsInfo.size()),
-                                      static_cast<int>(sizesInfo.size()),
-                                      static_cast<int>(stridesInfo.size())};
-  // If we parse an extra operand it needs to appear in the segmentSizes
-  if (parseExtraOperand)
-    segmentSizes.insert(segmentSizes.begin(), 1);
-  result.addAttribute(OpType::getOperandSegmentSizeAttr(),
-                      b.getI32VectorAttr(segmentSizes));
-
-  return failure(
-      parser.parseOptionalAttrDict(result.attributes) ||
+  Type srcType, dstType;
+  auto preResolutionFn = [&](OpAsmParser &parser, OperationState &result) {
+    return failure(parser.parseOptionalAttrDict(result.attributes) ||
                    parser.parseColonType(srcType) ||
-      parser.parseKeywordType(resultTypeKeyword.str().c_str(), dstType) ||
-      parser.resolveOperand(srcInfo, srcType, result.operands) ||
-      (parseExtraOperand &&
-       parser.resolveOperand(dstInfo, dstType, result.operands)) ||
-      parser.resolveOperands(offsetsInfo, indexType, result.operands) ||
-      parser.resolveOperands(sizesInfo, indexType, result.operands) ||
-      parser.resolveOperands(stridesInfo, indexType, result.operands) ||
-      parser.addTypeToList(dstType, result.types));
-}
-
-static ParseResult parseSubViewOp(OpAsmParser &parser, OperationState &result) {
-  return parseOpWithOffsetsSizesAndStrides<SubViewOp>(parser, result);
+                   parser.parseKeywordType("to", dstType) ||
+                   parser.resolveOperand(srcInfo, srcType, result.operands));
+  };
+  SmallVector<int, 4> segmentSizes{1}; // source memref
+  if (failed(parseOffsetsSizesAndStrides(parser, result, segmentSizes,
+                                         preResolutionFn)))
+    return failure();
+  return parser.addTypeToList(dstType, result.types);
 }
 
 void mlir::SubViewOp::build(OpBuilder &b, OperationState &result, Value source,
                             ArrayRef<int64_t> staticOffsets,
                             ArrayRef<int64_t> staticSizes,
                             ArrayRef<int64_t> staticStrides, ValueRange offsets,
                             ValueRange sizes, ValueRange strides,
                             ArrayRef<NamedAttribute> attrs) {
   auto sourceMemRefType = source.getType().cast<MemRefType>();
   auto resultType = inferResultType(sourceMemRefType, staticOffsets,
                                     staticSizes, staticStrides);
   build(b, result, resultType, source, offsets, sizes, strides,
         b.getI64ArrayAttr(staticOffsets), b.getI64ArrayAttr(staticSizes),
         b.getI64ArrayAttr(staticStrides));
   result.addAttributes(attrs);
 }
 
 /// Build a SubViewOp with all dynamic entries: `staticOffsets`, `staticSizes`
-/// and `staticStrides` are  automatically filled with source-memref-rank
+/// and `staticStrides` are automatically filled with source-memref-rank
 /// sentinel values that encode dynamic entries.
 void mlir::SubViewOp::build(OpBuilder &b, OperationState &result, Value source,
                             ValueRange offsets, ValueRange sizes,
                             ValueRange strides,
                             ArrayRef<NamedAttribute> attrs) {
   auto sourceMemRefType = source.getType().cast<MemRefType>();
   unsigned rank = sourceMemRefType.getRank();
   SmallVector<int64_t, 4> staticOffsetsVector;

 //===----------------------------------------------------------------------===//
 // SubTensorOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, SubTensorOp op) {
   return printOpWithOffsetsSizesAndStrides<SubTensorOp>(p, op);
 }
 
+/// Parse a subtensor op of the form:
+/// ```
+///   `subtensor` ssa-name
+///     `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
+///     `:` ranked-tensor-type `to` ranked-tensor-type
+/// ```
 static ParseResult parseSubTensorOp(OpAsmParser &parser,
                                     OperationState &result) {
-  return parseOpWithOffsetsSizesAndStrides<SubTensorOp>(parser, result);
+  OpAsmParser::OperandType srcInfo;
+  if (parser.parseOperand(srcInfo))
+    return failure();
+  Type srcType, dstType;
+  auto preResolutionFn = [&](OpAsmParser &parser, OperationState &result) {
+    return failure(parser.parseOptionalAttrDict(result.attributes) ||
+                   parser.parseColonType(srcType) ||
+                   parser.parseKeywordType("to", dstType) ||
+                   parser.resolveOperand(srcInfo, srcType, result.operands));
+  };
+  SmallVector<int, 4> segmentSizes{1}; // source tensor
+  if (failed(parseOffsetsSizesAndStrides(parser, result, segmentSizes,
+                                         preResolutionFn)))
+    return failure();
+  return parser.addTypeToList(dstType, result.types);
 }
 
 /// A subtensor result type can be fully inferred from the source type and the
 /// static representation of offsets, sizes and strides. Special sentinels
 /// encode the dynamic case.
 Type SubTensorOp::inferResultType(RankedTensorType sourceRankedTensorType,
                                   ArrayRef<int64_t> staticOffsets,
                                   ArrayRef<int64_t> staticSizes,

 
 static void print(OpAsmPrinter &p, SubTensorInsertOp op) {
   return printOpWithOffsetsSizesAndStrides<SubTensorInsertOp>(
       p, op,
       [](OpAsmPrinter &p, SubTensorInsertOp op) { p << " into " << op.dest(); },
       /*resultTypeKeyword=*/"into");
 }
 
+/// Parse a subtensor_insert op of the form:
+/// ```
+///   `subtensor_insert` ssa-name `into` ssa-name
+///     `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
+///     `:` ranked-tensor-type `into` ranked-tensor-type
+/// ```
 static ParseResult parseSubTensorInsertOp(OpAsmParser &parser,
                                           OperationState &result) {
-  return parseOpWithOffsetsSizesAndStrides<SubTensorInsertOp>(
-      parser, result,
-      [](OpAsmParser &parser, OpAsmParser::OperandType &dstInfo) {
-        return failure(parser.parseKeyword("into") ||
-                       parser.parseOperand(dstInfo));
-      },
-      "into");
+  OpAsmParser::OperandType srcInfo, dstInfo;
+  if (parser.parseOperand(srcInfo) || parser.parseKeyword("into") ||
+      parser.parseOperand(dstInfo))
+    return failure();
+  Type srcType, dstType;
+  auto preResolutionFn = [&](OpAsmParser &parser, OperationState &result) {
+    return failure(parser.parseOptionalAttrDict(result.attributes) ||
+                   parser.parseColonType(srcType) ||
+                   parser.parseKeywordType("into", dstType) ||
+                   parser.resolveOperand(srcInfo, srcType, result.operands) ||
+                   parser.resolveOperand(dstInfo, dstType, result.operands));
+  };
+  SmallVector<int, 4> segmentSizes{1, 1}; // source tensor, destination tensor
+  if (failed(parseOffsetsSizesAndStrides(parser, result, segmentSizes,
+                                         preResolutionFn)))
+    return failure();
+  return parser.addTypeToList(dstType, result.types);
 }
 
 void mlir::SubTensorInsertOp::build(
     OpBuilder &b, OperationState &result, Value source, Value dest,
     ArrayRef<int64_t> staticOffsets, ArrayRef<int64_t> staticSizes,
     ArrayRef<int64_t> staticStrides, ValueRange offsets, ValueRange sizes,
     ValueRange strides, ArrayRef<NamedAttribute> attrs) {
   build(b, result, dest.getType(), source, dest, offsets, sizes, strides,
         b.getI64ArrayAttr(staticOffsets), b.getI64ArrayAttr(staticSizes),
         b.getI64ArrayAttr(staticStrides));
   result.addAttributes(attrs);
 }
 
 /// Build a SubViewOp with all dynamic entries: `staticOffsets`, `staticSizes`
-/// and `staticStrides` are  automatically filled with source-memref-rank
+/// and `staticStrides` are automatically filled with source-memref-rank
 /// sentinel values that encode dynamic entries.
 void mlir::SubTensorInsertOp::build(OpBuilder &b, OperationState &result,
                                     Value source, Value dest,
                                     ValueRange offsets, ValueRange sizes,
                                     ValueRange strides,
                                     ArrayRef<NamedAttribute> attrs) {
   auto sourceRankedTensorType = source.getType().cast<RankedTensorType>();
   unsigned rank = sourceRankedTensorType.getRank();

mlir/lib/Interfaces/ViewLikeInterface.cpp

 //===- ViewLikeInterface.cpp - View-like operations in MLIR ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
-#include "mlir/IR/StandardTypes.h"
-
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ViewLike Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the loop-like interfaces.
 #include "mlir/Interfaces/ViewLikeInterface.cpp.inc"
LogicalResult mlir::verify(OffsetSizeAndStrideOpInterface op) {
   if (failed(verifyOpWithOffsetSizesAndStridesPart(
           op, "stride", ranks[2],
           OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
           op.static_strides(), ShapedType::isDynamicStrideOrOffset,
           op.strides())))
     return failure();
   return success();
 }
+
+/// Parse a mixed list with either (1) static integer values or (2) SSA values.
+/// Fill `result` with the integer ArrayAttr named `attrName` where `dynVal`
+/// encode the position of SSA values. Add the parsed SSA values to `ssa`
+/// in-order.
+//
+/// E.g. after parsing "[%arg0, 7, 42, %arg42]":
+///   1. `result` is filled with the i64 ArrayAttr "[`dynVal`, 7, 42, `dynVal`]"
+///   2. `ssa` is filled with "[%arg0, %arg1]".
+static ParseResult
+parseListOfOperandsOrIntegers(OpAsmParser &parser, OperationState &result,
+                              StringRef attrName, int64_t dynVal,
+                              SmallVectorImpl<OpAsmParser::OperandType> &ssa) {
+  if (failed(parser.parseLSquare()))
+    return failure();
+  // 0-D.
+  if (succeeded(parser.parseOptionalRSquare())) {
+    result.addAttribute(attrName, parser.getBuilder().getArrayAttr({}));
+    return success();
+  }
+
+  SmallVector<int64_t, 4> attrVals;
+  while (true) {
+    OpAsmParser::OperandType operand;
+    auto res = parser.parseOptionalOperand(operand);
+    if (res.hasValue() && succeeded(res.getValue())) {
+      ssa.push_back(operand);
+      attrVals.push_back(dynVal);
+    } else {
+      IntegerAttr attr;
+      if (failed(parser.parseAttribute<IntegerAttr>(attr)))
+        return parser.emitError(parser.getNameLoc())
+               << "expected SSA value or integer";
+      attrVals.push_back(attr.getInt());
+    }
+
+    if (succeeded(parser.parseOptionalComma()))
+      continue;
+    if (failed(parser.parseRSquare()))
+      return failure();
+    break;
+  }
+
+  auto arrayAttr = parser.getBuilder().getI64ArrayAttr(attrVals);
+  result.addAttribute(attrName, arrayAttr);
+  return success();
+}
+
+ParseResult mlir::parseOffsetsSizesAndStrides(
+    OpAsmParser &parser, OperationState &result, ArrayRef<int> segmentSizes,
+    llvm::function_ref<ParseResult(OpAsmParser &, OperationState &)>
+        preResolutionFn,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalOffsetPrefix,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalSizePrefix,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalStridePrefix) {
+  SmallVector<OpAsmParser::OperandType, 4> offsetsInfo, sizesInfo, stridesInfo;
+  auto indexType = parser.getBuilder().getIndexType();
+  if ((parseOptionalOffsetPrefix && parseOptionalOffsetPrefix(parser)) ||
+      parseListOfOperandsOrIntegers(
+          parser, result,
+          OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName(),
+          ShapedType::kDynamicStrideOrOffset, offsetsInfo) ||
+      (parseOptionalSizePrefix && parseOptionalSizePrefix(parser)) ||
+      parseListOfOperandsOrIntegers(
+          parser, result,
+          OffsetSizeAndStrideOpInterface::getStaticSizesAttrName(),
+          ShapedType::kDynamicSize, sizesInfo) ||
+      (parseOptionalStridePrefix && parseOptionalStridePrefix(parser)) ||
+      parseListOfOperandsOrIntegers(
+          parser, result,
+          OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
+          ShapedType::kDynamicStrideOrOffset, stridesInfo))
+    return failure();
+  // Add segment sizes to result
+  SmallVector<int, 4> segmentSizesFinal(segmentSizes.begin(),
+                                        segmentSizes.end());
+  segmentSizesFinal.append({static_cast<int>(offsetsInfo.size()),
+                            static_cast<int>(sizesInfo.size()),
+                            static_cast<int>(stridesInfo.size())});
+  auto b = parser.getBuilder();
+  result.addAttribute(
+      OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr(),
+      b.getI32VectorAttr(segmentSizesFinal));
+  return failure(
+      (preResolutionFn && preResolutionFn(parser, result)) ||
+      parser.resolveOperands(offsetsInfo, indexType, result.operands) ||
+      parser.resolveOperands(sizesInfo, indexType, result.operands) ||
+      parser.resolveOperands(stridesInfo, indexType, result.operands));
+}