[mlir] NFC - Refactor and expose a helper printOffsetSizesAndStrides helper function.

[nicolasvasilache]

Summary

Print part of an op of the form:

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

[nicolasvasilache]

Rebase

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

code extraBaseClassDeclaration = [{
     operand_range getDynamicSizes() { return sizes(); }
 
     /// Return the list of Range (i.e. offset, size, stride). Each
     /// Range entry contains either the dynamic value or a ConstantIndexOp
     /// constructed with `b` at location `loc`.
     SmallVector<Range, 8> getOrCreateRanges(OpBuilder &b, Location loc) {
       return mlir::getOrCreateRanges(*this, b, loc);
     }
-
-    static ArrayRef<StringRef> getSpecialAttrNames() {
-      static SmallVector<StringRef, 4> names{
-        OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName(),
-        OffsetSizeAndStrideOpInterface::getStaticSizesAttrName(),
-        OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
-        getOperandSegmentSizeAttr()};
-      return names;
-   }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // AbsFOp
 //===----------------------------------------------------------------------===//
 
 def AbsFOp : FloatUnaryOp<"absf"> {

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 struct Range {
   Value offset;
   Value size;
   Value stride;
 };
 
 class OffsetSizeAndStrideOpInterface;
 LogicalResult verify(OffsetSizeAndStrideOpInterface op);
+} // namespace mlir
+
+/// Include the generated interface declarations.
+#include "mlir/Interfaces/ViewLikeInterface.h.inc"
+
+namespace mlir {
+/// Print part of an op of the form:
+/// ```
+///   <optional-offset-prefix>`[` offset-list `]`
+///   <optional-size-prefix>`[` size-list `]`
+///   <optional-stride-prefix>[` stride-list `]`
+/// ```
+void printOffsetsSizesAndStrides(
+    OpAsmPrinter &p, OffsetSizeAndStrideOpInterface op,
+    StringRef offsetPrefix = "", StringRef sizePrefix = " ",
+    StringRef stridePrefix = " ",
+    ArrayRef<StringRef> elidedAttrs =
+        OffsetSizeAndStrideOpInterface::getSpecialAttrNames());
 
 /// 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
ParseResult parseOffsetsSizesAndStrides(
     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);
+/// `preResolutionFn`-less version of `parseOffsetsSizesAndStrides`.
+ParseResult parseOffsetsSizesAndStrides(
+    OpAsmParser &parser, OperationState &result, ArrayRef<int> segmentSizes,
+    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/include/mlir/Interfaces/ViewLikeInterface.td

static StringRef getStaticOffsetsAttrName() {
       return "static_offsets";
     }
     static StringRef getStaticSizesAttrName() {
       return "static_sizes";
     }
     static StringRef getStaticStridesAttrName() {
       return "static_strides";
     }
+    static ArrayRef<StringRef> getSpecialAttrNames() {
+      static SmallVector<StringRef, 4> names{
+        OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName(),
+        OffsetSizeAndStrideOpInterface::getStaticSizesAttrName(),
+        OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
+        OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr()};
+      return names;
+   }
   }];
 
   let verify = [{
     return mlir::verify(cast<OffsetSizeAndStrideOpInterface>($_op));
   }];
 }
 
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE

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

for (auto operandIndex :
 
       fusableDependences[producerOp.getOperation()].push_back(
           *fusableDependence);
     }
   }
   return fusableDependences;
 }
 
-static bool isZero(Value v) {
-  if (auto cst = v.getDefiningOp<ConstantIndexOp>())
-    return cst.getValue() == 0;
-  return false;
-}
-
 /// Tile the fused loops in the root operation, by setting the tile sizes for
 /// all other loops to zero (those will be tiled later).
 static Optional<TiledLinalgOp> tileRootOperation(
     OpBuilder &builder, LinalgOp op, ArrayRef<Value> tileSizeVector,
     const LinalgTilingOptions &options, const std::set<unsigned> &fusedLoops) {
   SmallVector<Value, 4> tileSizes(tileSizeVector.begin(), tileSizeVector.end());
   auto zero = std_constant_index(0);
   for (unsigned i = 0, e = tileSizes.size(); i != e; ++i)

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

OpFoldResult AddIOp::fold(ArrayRef<Attribute> operands) {
   /// addi(x, 0) -> x
   if (matchPattern(rhs(), m_Zero()))
     return lhs();
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a + b; });
 }
 
-//===----------------------------------------------------------------------===//
-// BaseOpWithOffsetSizesAndStridesOp
-//===----------------------------------------------------------------------===//
-
-/// Print a list with either (1) the static integer value in `arrayAttr` if
-/// `isDynamic` evaluates to false or (2) the next value otherwise.
-/// This allows idiomatic printing of mixed value and integer attributes in a
-/// list. E.g. `[%arg0, 7, 42, %arg42]`.
-static void
-printListOfOperandsOrIntegers(OpAsmPrinter &p, ValueRange values,
-                              ArrayAttr arrayAttr,
-                              llvm::function_ref<bool(int64_t)> isDynamic) {
-  p << '[';
-  unsigned idx = 0;
-  llvm::interleaveComma(arrayAttr, p, [&](Attribute a) {
-    int64_t val = a.cast<IntegerAttr>().getInt();
-    if (isDynamic(val))
-      p << values[idx++];
-    else
-      p << val;
-  });
-  p << ']';
-}
-
-/// 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 ")
-           << expectedNumElements << " " << name << " values";
-  unsigned expectedNumDynamicEntries =
-      llvm::count_if(attr.getValue(), [&](Attribute attr) {
-        return isDynamic(attr.cast<IntegerAttr>().getInt());
-      });
-  if (values.size() != expectedNumDynamicEntries)
-    return op.emitError("expected ")
-           << expectedNumDynamicEntries << " dynamic " << name << " values";
-  return success();
-}
-
 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
   return llvm::to_vector<4>(
       llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
         return a.cast<IntegerAttr>().getInt();
       }));
 }
 
void mlir::MemRefReinterpretCastOp::build(OpBuilder &b, OperationState &result,
   SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamicSize);
   SmallVector<int64_t, 4> staticStridesVector(
       rank, ShapedType::kDynamicStrideOrOffset);
   build(b, result, resultType, source,
         /*staticOffset=*/ShapedType::kDynamicStrideOrOffset, staticSizesVector,
         staticStridesVector, offset, sizes, strides, attrs);
 }
 
-/// Print of the form:
+/// Print 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 void print(OpAsmPrinter &p, MemRefReinterpretCastOp op) {
   int stdDotLen = StandardOpsDialect::getDialectNamespace().size() + 1;
-  p << op.getOperationName().drop_front(stdDotLen) << " " << op.source()
-    << " to offset: ";
-  printListOfOperandsOrIntegers(p, op.offsets(), op.static_offsets(),
-                                ShapedType::isDynamicStrideOrOffset);
-  p << ", sizes: ";
-  printListOfOperandsOrIntegers(p, op.sizes(), op.static_sizes(),
-                                ShapedType::isDynamic);
-  p << ", strides: ";
-  printListOfOperandsOrIntegers(p, op.strides(), op.static_strides(),
-                                ShapedType::isDynamicStrideOrOffset);
-  p.printOptionalAttrDict(
-      op.getAttrs(),
-      /*elidedAttrs=*/MemRefReinterpretCastOp::getSpecialAttrNames());
+  p << op.getOperation()->getName().getStringRef().drop_front(stdDotLen) << ' ';
+  p << op.source() << " ";
+  printOffsetsSizesAndStrides(
+      p, op, /*offsetPrefix=*/"to offset: ", /*sizePrefix=*/", sizes: ",
+      /*stridePrefix=*/", strides: ");
   p << ": " << op.source().getType() << " to " << op.getType();
 }
 
 /// Parse a memref_reinterpret_cast op of the form:
 /// ```
 ///   `memref_reinterpret_cast` ssa-name to
 ///       offset: `[` offset `]`
 ///       sizes: `[` size-list `]`
static ParseResult parseMemRefReinterpretCastOp(OpAsmParser &parser,
 
   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,
+  if (failed(parseOffsetsSizesAndStrides(parser, result,
+                                         /*segmentSizes=*/{1}, // source memref
                                          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.
Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
   // The type is now known.
   return MemRefType::get(
       staticSizes, sourceMemRefType.getElementType(),
       makeStridedLinearLayoutMap(targetStrides, targetOffset,
                                  sourceMemRefType.getContext()),
       sourceMemRefType.getMemorySpace());
 }
 
-/// Print SubViewOp in the form:
+/// Print a subview op of the form:
 /// ```
-///   subview ssa-name `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
+///   `subview` ssa-name
+///     `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
 ///     `:` strided-memref-type `to` strided-memref-type
 /// ```
-template <typename OpType>
-static void printOpWithOffsetsSizesAndStrides(
-    OpAsmPrinter &p, OpType op,
-    llvm::function_ref<void(OpAsmPrinter &p, OpType op)> printExtraOperands =
-        [](OpAsmPrinter &p, OpType op) {},
-    StringRef resultTypeKeyword = "to") {
+static void print(OpAsmPrinter &p, SubViewOp op) {
   int stdDotLen = StandardOpsDialect::getDialectNamespace().size() + 1;
   p << op.getOperation()->getName().getStringRef().drop_front(stdDotLen) << ' ';
   p << op.source();
-  printExtraOperands(p, op);
-  printListOfOperandsOrIntegers(p, op.offsets(), op.static_offsets(),
-                                ShapedType::isDynamicStrideOrOffset);
-  p << ' ';
-  printListOfOperandsOrIntegers(p, op.sizes(), op.static_sizes(),
-                                ShapedType::isDynamic);
-  p << ' ';
-  printListOfOperandsOrIntegers(p, op.strides(), op.static_strides(),
-                                ShapedType::isDynamicStrideOrOffset);
-  p << ' ';
-  p.printOptionalAttrDict(op.getAttrs(),
-                          /*elidedAttrs=*/{OpType::getSpecialAttrNames()});
-  p << " : " << op.getSourceType() << " " << resultTypeKeyword << " "
-    << op.getType();
-}
-
-static void print(OpAsmPrinter &p, SubViewOp op) {
-  return printOpWithOffsetsSizesAndStrides<SubViewOp>(p, op);
+  printOffsetsSizesAndStrides(p, op);
+  p << " : " << op.getSourceType() << " to " << op.getType();
 }
 
 /// Parse a subview op of the form:
 /// ```
 ///   `subview` ssa-name
 ///     `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
 ///     `:` strided-memref-type `to` strided-memref-type
 /// ```
 static ParseResult parseSubViewOp(OpAsmParser &parser, OperationState &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 memref
-  if (failed(parseOffsetsSizesAndStrides(parser, result, segmentSizes,
+
+  if (failed(parseOffsetsSizesAndStrides(parser, result,
+                                         /*segmentSizes=*/{1}, // source memref
                                          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,
OpFoldResult SubViewOp::fold(ArrayRef<Attribute> operands) {
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // SubTensorOp
 //===----------------------------------------------------------------------===//
 
+/// Print a subtensor op of the form:
+/// ```
+///   `subtensor` ssa-name
+///     `[` offset-list `]` `[` size-list `]` `[` stride-list `]`
+///     `:` ranked-tensor-type `to` ranked-tensor-type
+/// ```
 static void print(OpAsmPrinter &p, SubTensorOp op) {
-  return printOpWithOffsetsSizesAndStrides<SubTensorOp>(p, op);
+  int stdDotLen = StandardOpsDialect::getDialectNamespace().size() + 1;
+  p << op.getOperation()->getName().getStringRef().drop_front(stdDotLen) << ' ';
+  p << op.source();
+  printOffsetsSizesAndStrides(p, op);
+  p << " : " << op.getSourceType() << " to " << op.getType();
 }
 
 /// 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) {
   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,
+
+  if (failed(parseOffsetsSizesAndStrides(parser, result,
+                                         /*segmentSizes=*/{1}, // source tensor
                                          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.
results
       .insert<OpWithOffsetSizesAndStridesConstantArgumentFolder<SubTensorOp>>(
           context);
 }
 
 //===----------------------------------------------------------------------===//
 // SubTensorInsertOp
 //===----------------------------------------------------------------------===//
 
+/// Print 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 void print(OpAsmPrinter &p, SubTensorInsertOp op) {
-  return printOpWithOffsetsSizesAndStrides<SubTensorInsertOp>(
-      p, op,
-      [](OpAsmPrinter &p, SubTensorInsertOp op) { p << " into " << op.dest(); },
-      /*resultTypeKeyword=*/"into");
+  int stdDotLen = StandardOpsDialect::getDialectNamespace().size() + 1;
+  p << op.getOperation()->getName().getStringRef().drop_front(stdDotLen) << ' ';
+  p << op.source() << " into " << op.dest();
+  printOffsetsSizesAndStrides(p, op);
+  p << " : " << op.getSourceType() << " into " << op.getType();
 }
 
 /// 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) {
   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)))
+
+  if (failed(parseOffsetsSizesAndStrides(
+          parser, result,
+          /*segmentSizes=*/{1, 1}, // source tensor, destination tensor
     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,

mlir/lib/Interfaces/ViewLikeInterface.cpp

if (failed(verifyOpWithOffsetSizesAndStridesPart(
           op, "stride", ranks[2],
           OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
           op.static_strides(), ShapedType::isDynamicStrideOrOffset,
           op.strides())))
     return failure();
   return success();
 }
 
+/// Print a list with either (1) the static integer value in `arrayAttr` if
+/// `isDynamic` evaluates to false or (2) the next value otherwise.
+/// This allows idiomatic printing of mixed value and integer attributes in a
+/// list. E.g. `[%arg0, 7, 42, %arg42]`.
+static void
+printListOfOperandsOrIntegers(OpAsmPrinter &p, ValueRange values,
+                              ArrayAttr arrayAttr,
+                              llvm::function_ref<bool(int64_t)> isDynamic) {
+  p << '[';
+  unsigned idx = 0;
+  llvm::interleaveComma(arrayAttr, p, [&](Attribute a) {
+    int64_t val = a.cast<IntegerAttr>().getInt();
+    if (isDynamic(val))
+      p << values[idx++];
+    else
+      p << val;
+  });
+  p << ']';
+}
+
+void mlir::printOffsetsSizesAndStrides(OpAsmPrinter &p,
+                                       OffsetSizeAndStrideOpInterface op,
+                                       StringRef offsetPrefix,
+                                       StringRef sizePrefix,
+                                       StringRef stridePrefix,
+                                       ArrayRef<StringRef> elidedAttrs) {
+  p << offsetPrefix;
+  printListOfOperandsOrIntegers(p, op.offsets(), op.static_offsets(),
+                                ShapedType::isDynamicStrideOrOffset);
+  p << sizePrefix;
+  printListOfOperandsOrIntegers(p, op.sizes(), op.static_sizes(),
+                                ShapedType::isDynamic);
+  p << stridePrefix;
+  printListOfOperandsOrIntegers(p, op.strides(), op.static_strides(),
+                                ShapedType::isDynamicStrideOrOffset);
+  p.printOptionalAttrDict(op.getAttrs(), elidedAttrs);
+}
+
 /// 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]".
parseListOfOperandsOrIntegers(OpAsmParser &parser, OperationState &result,
   }
 
   auto arrayAttr = parser.getBuilder().getI64ArrayAttr(attrVals);
   result.addAttribute(attrName, arrayAttr);
   return success();
 }
 
 ParseResult mlir::parseOffsetsSizesAndStrides(
-    OpAsmParser &parser,
-    OperationState &result,
-    ArrayRef<int> segmentSizes,
+    OpAsmParser &parser, OperationState &result, ArrayRef<int> segmentSizes,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalOffsetPrefix,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalSizePrefix,
+    llvm::function_ref<ParseResult(OpAsmParser &)> parseOptionalStridePrefix) {
+  return parseOffsetsSizesAndStrides(
+      parser, result, segmentSizes, nullptr, parseOptionalOffsetPrefix,
+      parseOptionalSizePrefix, parseOptionalStridePrefix);
+}
+
+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());
+  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));
+      parser.getBuilder().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));
 }