[MLIR] Make the implementations for getMixedOffsets/Sizes/Strides independent of OffsetSizeAndStrideOpInterface

[frgossen]

Summary

The functions are effectively independent of the interface already, however, they take it as an argument for no reason.
The current state complicates reuse outside of MLIR.

[frgossen]

Rebase

[frgossen]

Rebase

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpImplementation.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 {
   OpFoldResult offset;
   OpFoldResult size;
   OpFoldResult stride;
 };
 
-class OffsetSizeAndStrideOpInterface;
-
-/// Return a vector of all the static or dynamic offsets of the op from provided
-/// external static and dynamic offsets.
-SmallVector<OpFoldResult, 4> getMixedOffsets(OffsetSizeAndStrideOpInterface op,
-                                             ArrayAttr staticOffsets,
-                                             ValueRange offsets);
-
-/// Return a vector of all the static or dynamic sizes of the op from provided
-/// external static and dynamic sizes.
-SmallVector<OpFoldResult, 4> getMixedSizes(OffsetSizeAndStrideOpInterface op,
-                                           ArrayAttr staticSizes,
-                                           ValueRange sizes);
-
-/// Return a vector of all the static or dynamic strides of the op from provided
-/// external static and dynamic strides.
-SmallVector<OpFoldResult, 4> getMixedStrides(OffsetSizeAndStrideOpInterface op,
-                                             ArrayAttr staticStrides,
-                                             ValueRange strides);
+/// Return a vector of all the static and dynamic offsets/strides.
+SmallVector<OpFoldResult, 4> getMixedStridesOrOffsets(ArrayAttr staticValues,
+                                                      ValueRange dynamicValues);
+
+/// Return a vector of all the static and dynamic sizes.
+SmallVector<OpFoldResult, 4> getMixedSizes(ArrayAttr staticValues,
+                                           ValueRange dynamicValues);
 
-/// Decompose a vector of mixed static or dynamic strides/offsets into the
+/// Decompose a vector of mixed static and dynamic strides/offsets into the
 /// corresponding pair of arrays. This is the inverse function of
-/// `getMixedStrides` and `getMixedOffsets`.
+/// `getMixedStridesOrOffsets`.
 std::pair<ArrayAttr, SmallVector<Value>> decomposeMixedStridesOrOffsets(
     OpBuilder &b, const SmallVectorImpl<OpFoldResult> &mixedValues);
 
 /// Decompose a vector of mixed static or dynamic strides/offsets into the
 /// corresponding pair of arrays. This is the inverse function of
 /// `getMixedSizes`.
 std::pair<ArrayAttr, SmallVector<Value>>
 decomposeMixedSizes(OpBuilder &b,
                     const SmallVectorImpl<OpFoldResult> &mixedValues);
 
+class OffsetSizeAndStrideOpInterface;
+
 namespace detail {
+
 LogicalResult verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op);
 
 bool sameOffsetsSizesAndStrides(
     OffsetSizeAndStrideOpInterface a, OffsetSizeAndStrideOpInterface b,
     llvm::function_ref<bool(OpFoldResult, OpFoldResult)> cmp);
+
 } // namespace detail
 } // namespace mlir
 
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/ViewLikeInterface.h.inc"
 
 namespace mlir {
 

mlir/include/mlir/Interfaces/ViewLikeInterface.td

InterfaceMethod<
       /*desc=*/[{
         Return a vector of all the static or dynamic sizes of the op.
       }],
       /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult, 4>",
       /*methodName=*/"getMixedOffsets",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return ::mlir::getMixedOffsets($_op, $_op.static_offsets(),
+        return ::mlir::getMixedStridesOrOffsets($_op.static_offsets(),
                                                 $_op.offsets());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return a vector of all the static or dynamic sizes of the op.
       }],
       /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult, 4>",
       /*methodName=*/"getMixedSizes",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return ::mlir::getMixedSizes($_op, $_op.static_sizes(), $_op.sizes());
+        return ::mlir::getMixedSizes($_op.static_sizes(), $_op.sizes());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return a vector of all the static or dynamic strides of the op.
       }],
       /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult, 4>",
       /*methodName=*/"getMixedStrides",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return ::mlir::getMixedStrides($_op, $_op.static_strides(),
+        return ::mlir::getMixedStridesOrOffsets($_op.static_strides(),
                                                 $_op.strides());
       }]
     >,
 
     InterfaceMethod<
       /*desc=*/[{
         Return true if the offset `idx` is dynamic.
       }],
       /*retTy=*/"bool",

mlir/lib/Interfaces/ViewLikeInterface.cpp

for (auto it : llvm::zip(a.getMixedSizes(), b.getMixedSizes()))
     if (!cmp(std::get<0>(it), std::get<1>(it)))
       return false;
   for (auto it : llvm::zip(a.getMixedStrides(), b.getMixedStrides()))
     if (!cmp(std::get<0>(it), std::get<1>(it)))
       return false;
   return true;
 }
 
-SmallVector<OpFoldResult, 4>
-mlir::getMixedOffsets(OffsetSizeAndStrideOpInterface op,
-                      ArrayAttr staticOffsets, ValueRange offsets) {
-  SmallVector<OpFoldResult, 4> res;
-  unsigned numDynamic = 0;
-  unsigned count = static_cast<unsigned>(staticOffsets.size());
-  for (unsigned idx = 0; idx < count; ++idx) {
-    if (op.isDynamicOffset(idx))
-      res.push_back(offsets[numDynamic++]);
-    else
-      res.push_back(staticOffsets[idx]);
-  }
-  return res;
+static SmallVector<OpFoldResult>
+getMixedImpl(ArrayAttr staticValues, ValueRange dynamicValues,
+             const int64_t dynamicValuePlaceholder) {
+  int64_t idxDynamic = 0;
+  SmallVector<OpFoldResult> result;
+  for (const Attribute &staticAttr : staticValues) {
+    int64_t staticInt = staticAttr.cast<IntegerAttr>().getInt();
+    if (staticInt == dynamicValuePlaceholder) {
+      result.push_back(dynamicValues[idxDynamic++]);
+    } else {
+      result.push_back(staticAttr);
     }
-
-SmallVector<OpFoldResult, 4>
-mlir::getMixedSizes(OffsetSizeAndStrideOpInterface op, ArrayAttr staticSizes,
-                    ValueRange sizes) {
-  SmallVector<OpFoldResult, 4> res;
-  unsigned numDynamic = 0;
-  unsigned count = static_cast<unsigned>(staticSizes.size());
-  for (unsigned idx = 0; idx < count; ++idx) {
-    if (op.isDynamicSize(idx))
-      res.push_back(sizes[numDynamic++]);
-    else
-      res.push_back(staticSizes[idx]);
   }
-  return res;
+  return result;
 }
 
 SmallVector<OpFoldResult, 4>
-mlir::getMixedStrides(OffsetSizeAndStrideOpInterface op,
-                      ArrayAttr staticStrides, ValueRange strides) {
-  SmallVector<OpFoldResult, 4> res;
-  unsigned numDynamic = 0;
-  unsigned count = static_cast<unsigned>(staticStrides.size());
-  for (unsigned idx = 0; idx < count; ++idx) {
-    if (op.isDynamicStride(idx))
-      res.push_back(strides[numDynamic++]);
-    else
-      res.push_back(staticStrides[idx]);
-  }
-  return res;
+mlir::getMixedStridesOrOffsets(ArrayAttr staticValues,
+                               ValueRange dynamicValues) {
+  return getMixedImpl(staticValues, dynamicValues,
+                      ShapedType::kDynamicStrideOrOffset);
+
+SmallVector<OpFoldResult, 4> mlir::getMixedSizes(ArrayAttr staticValues,
+                                                 ValueRange dynamicValues) {
+  return getMixedImpl(staticValues, dynamicValues, ShapedType::kDynamicSize);
 }
 
 static std::pair<ArrayAttr, SmallVector<Value>>
 decomposeMixedImpl(OpBuilder &b,
                    const SmallVectorImpl<OpFoldResult> &mixedValues,
                    const int64_t dynamicValuePlaceholder) {
   SmallVector<int64_t> staticValues;
   SmallVector<Value> dynamicValues;