[mlir][NFC] Add a new getStridesAndOffset function

[qcolombet]

Summary

The new function is a wrapper around the regular getStridesAndOffset that offers a more compact way (as in writing less code) of getting the relevant information.

This method is intended to be used only when it is known that the LogicalResult of the regular getStridesAndOffset must be "succeeded".

This wrapper will assert on that.

Note: I feel we may want to give this function a different name so that it is explicit that it should be used only when we know it is valid. I couldn't come up with a good name though. What do people think?

[jreiffers]

Thanks for the cleanup!

mlir/include/mlir/IR/BuiltinTypes.h

 /// A stride specification is a list of integer values that are either static
 /// or dynamic (encoded with ShapedType::kDynamic). Strides encode
 /// the distance in the number of elements between successive entries along a
 /// particular dimension.
 LogicalResult getStridesAndOffset(MemRefType t,
                                   SmallVectorImpl<int64_t> &strides,
                                   int64_t &offset);
 
+/// Wrapper around getStridesAndOffset(MemRefType, SmallVectorImpl<int64_t>,
+/// int64_t) that will assert if the logical result is not succeeded.
+std::pair<SmallVector<int64_t>, int64_t> getStridesAndOffset(MemRefType t);
+
 /// Return a version of `t` with identity layout if it can be determined
 /// statically that the layout is the canonical contiguous strided layout.
 /// Otherwise pass `t`'s layout into `simplifyAffineMap` and return a copy of
 /// `t` with simplified layout.
 MemRefType canonicalizeStridedLayout(MemRefType t);
 
 /// Given MemRef `sizes` that are either static or dynamic, returns the
 /// canonical "contiguous" strides AffineExpr. Strides are multiplicative and

mlir/lib/Conversion/LLVMCommon/MemRefBuilder.cpp

 }
 
 MemRefDescriptor MemRefDescriptor::fromStaticShape(
     OpBuilder &builder, Location loc, LLVMTypeConverter &typeConverter,
     MemRefType type, Value memory, Value alignedMemory) {
   assert(type.hasStaticShape() && "unexpected dynamic shape");
 
   // Extract all strides and offsets and verify they are static.
-  int64_t offset;
-  SmallVector<int64_t, 4> strides;
-  auto result = getStridesAndOffset(type, strides, offset);
-  (void)result;
-  assert(succeeded(result) && "unexpected failure in stride computation");
+  auto [strides, offset] = getStridesAndOffset(type);
   assert(!ShapedType::isDynamic(offset) &&
          "expected static offset");
   assert(!llvm::any_of(strides, ShapedType::isDynamic) &&
          "expected static strides");
 
   auto convertedType = typeConverter.convertType(type);
   assert(convertedType && "unexpected failure in memref type conversion");
 

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

Value ConvertToLLVMPattern::createIndexConstant(
     ConversionPatternRewriter &builder, Location loc, uint64_t value) const {
   return createIndexAttrConstant(builder, loc, getIndexType(), value);
 }
 
 Value ConvertToLLVMPattern::getStridedElementPtr(
     Location loc, MemRefType type, Value memRefDesc, ValueRange indices,
     ConversionPatternRewriter &rewriter) const {
 
-  int64_t offset;
-  SmallVector<int64_t, 4> strides;
-  auto successStrides = getStridesAndOffset(type, strides, offset);
-  assert(succeeded(successStrides) && "unexpected non-strided memref");
-  (void)successStrides;
+  auto [strides, offset] = getStridesAndOffset(type);
 
   MemRefDescriptor memRefDescriptor(memRefDesc);
   Value base = memRefDescriptor.alignedPtr(rewriter, loc);
 
   Value index;
   if (offset != 0) // Skip if offset is zero.
     index = ShapedType::isDynamic(offset)
                 ? memRefDescriptor.offset(rewriter, loc)

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
                                 ArrayRef<int64_t> staticStrides) {
   unsigned rank = sourceMemRefType.getRank();
   (void)rank;
   assert(staticOffsets.size() == rank && "staticOffsets length mismatch");
   assert(staticSizes.size() == rank && "staticSizes length mismatch");
   assert(staticStrides.size() == rank && "staticStrides length mismatch");
 
   // Extract source offset and strides.
-  int64_t sourceOffset;
-  SmallVector<int64_t, 4> sourceStrides;
-  auto res = getStridesAndOffset(sourceMemRefType, sourceStrides, sourceOffset);
-  assert(succeeded(res) && "SubViewOp expected strided memref type");
-  (void)res;
+  auto [sourceStrides, sourceOffset] = getStridesAndOffset(sourceMemRefType);
 
   // Compute target offset whose value is:
   //   `sourceOffset + sum_i(staticOffset_i * sourceStrides_i)`.
   int64_t targetOffset = sourceOffset;
   for (auto it : llvm::zip(staticOffsets, sourceStrides)) {
     auto staticOffset = std::get<0>(it), targetStride = std::get<1>(it);
     using saturated_arith::Wrapper;
     targetOffset =
void TransposeOp::getAsmResultNames(
   setNameFn(getResult(), "transpose");
 }
 
 /// Build a strided memref type by applying `permutationMap` tp `memRefType`.
 static MemRefType inferTransposeResultType(MemRefType memRefType,
                                            AffineMap permutationMap) {
   auto rank = memRefType.getRank();
   auto originalSizes = memRefType.getShape();
-  int64_t offset;
-  SmallVector<int64_t, 4> originalStrides;
-  auto res = getStridesAndOffset(memRefType, originalStrides, offset);
-  assert(succeeded(res) &&
-         originalStrides.size() == static_cast<unsigned>(rank));
-  (void)res;
+  auto [originalStrides, offset] = getStridesAndOffset(memRefType);
+  assert(originalStrides.size() == static_cast<unsigned>(rank));
 
   // Compute permuted sizes and strides.
   SmallVector<int64_t> sizes(rank, 0);
   SmallVector<int64_t> strides(rank, 1);
   for (const auto &en : llvm::enumerate(permutationMap.getResults())) {
     unsigned position = en.value().cast<AffineDimExpr>().getPosition();
     sizes[en.index()] = originalSizes[position];
     strides[en.index()] = originalStrides[position];

mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp

LogicalResult matchAndRewrite(memref::SubViewOp subview,
     Location origLoc = subview.getLoc();
     Value source = subview.getSource();
     auto sourceType = source.getType().cast<MemRefType>();
     unsigned sourceRank = sourceType.getRank();
 
     auto newExtractStridedMetadata =
         rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
 
-    SmallVector<int64_t> sourceStrides;
-    int64_t sourceOffset;
-
-    bool hasKnownStridesAndOffset =
-        succeeded(getStridesAndOffset(sourceType, sourceStrides, sourceOffset));
-    (void)hasKnownStridesAndOffset;
-    assert(hasKnownStridesAndOffset &&
-           "getStridesAndOffset must work on valid subviews");
+    auto [sourceStrides, sourceOffset] = getStridesAndOffset(sourceType);
 
     // Compute the new strides and offset from the base strides and offset:
     // newStride#i = baseStride#i * subStride#i
     // offset = baseOffset + sum(subOffsets#i * newStrides#i)
     SmallVector<OpFoldResult> strides;
     SmallVector<OpFoldResult> subStrides = subview.getMixedStrides();
     auto origStrides = newExtractStridedMetadata.getStrides();
 
for (int i = groupSize - 1; i >= 0; --i) {
     }
 
     currentStride *= dimSize;
   }
 
   // Collect the statically known information about the original stride.
   Value source = expandShape.getSrc();
   auto sourceType = source.getType().cast<MemRefType>();
-  SmallVector<int64_t> strides;
-  int64_t offset;
-  bool hasKnownStridesAndOffset =
-      succeeded(getStridesAndOffset(sourceType, strides, offset));
-  (void)hasKnownStridesAndOffset;
-  assert(hasKnownStridesAndOffset &&
-         "getStridesAndOffset must work on valid expand_shape");
+  auto [strides, offset] = getStridesAndOffset(sourceType);
 
   OpFoldResult origStride =
       ShapedType::isDynamic(strides[groupId])
           ? origStrides[groupId]
           : builder.getIndexAttr(strides[groupId]);
 
   // Apply the original stride to all the strides.
   int64_t doneStrideIdx = 0;
getCollapsedStride(memref::CollapseShapeOp collapseShape, OpBuilder &builder,
   SmallVector<int64_t, 2> reassocGroup =
       collapseShape.getReassociationIndices()[groupId];
   assert(!reassocGroup.empty() &&
          "Reassociation group should have at least one dimension");
 
   Value source = collapseShape.getSrc();
   auto sourceType = source.getType().cast<MemRefType>();
 
-  SmallVector<int64_t> strides;
-  int64_t offset;
-  bool hasKnownStridesAndOffset =
-      succeeded(getStridesAndOffset(sourceType, strides, offset));
-  (void)hasKnownStridesAndOffset;
-  assert(hasKnownStridesAndOffset &&
-         "getStridesAndOffset must work on valid collapse_shape");
+  auto [strides, offset] = getStridesAndOffset(sourceType);
 
   SmallVector<OpFoldResult> collapsedStride;
   int64_t innerMostDimForGroup = reassocGroup.back();
   int64_t innerMostStrideForGroup = strides[innerMostDimForGroup];
   collapsedStride.push_back(
       ShapedType::isDynamic(innerMostStrideForGroup)
           ? origStrides[innerMostDimForGroup]
           : builder.getIndexAttr(innerMostStrideForGroup));
LogicalResult matchAndRewrite(ReassociativeReshapeLikeOp reshape,
     Value source = reshape.getSrc();
     auto sourceType = source.getType().cast<MemRefType>();
     unsigned sourceRank = sourceType.getRank();
 
     auto newExtractStridedMetadata =
         rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
 
     // Collect statically known information.
-    SmallVector<int64_t> strides;
-    int64_t offset;
-    bool hasKnownStridesAndOffset =
-        succeeded(getStridesAndOffset(sourceType, strides, offset));
-    (void)hasKnownStridesAndOffset;
-    assert(hasKnownStridesAndOffset &&
-           "getStridesAndOffset must work on valid reassociative_reshape_like");
+    auto [strides, offset] = getStridesAndOffset(sourceType);
     MemRefType reshapeType = reshape.getResultType();
     unsigned reshapeRank = reshapeType.getRank();
 
     OpFoldResult offsetOfr =
         ShapedType::isDynamic(offset)
             ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
             : rewriter.getIndexAttr(offset);
 

mlir/lib/IR/BuiltinTypes.cpp

for (auto e : strideExprs) {
     if (auto c = e.dyn_cast<AffineConstantExpr>())
       strides.push_back(c.getValue());
     else
       strides.push_back(ShapedType::kDynamic);
   }
   return success();
 }
 
+std::pair<SmallVector<int64_t>, int64_t>
+mlir::getStridesAndOffset(MemRefType t) {
+  SmallVector<int64_t> strides;
+  int64_t offset;
+  LogicalResult status = getStridesAndOffset(t, strides, offset);
+  (void)status;
+  assert(succeeded(status) && "Invalid use of check-free getStridesAndOffset");
+  return {strides, offset};
+}
+
 //===----------------------------------------------------------------------===//
 /// TupleType
 //===----------------------------------------------------------------------===//
 
 /// Return the elements types for this tuple.
 ArrayRef<Type> TupleType::getTypes() const { return getImpl()->getTypes(); }
 
 /// Accumulate the types contained in this tuple and tuples nested within it.