[Mlir-commits] [mlir] 02d29af - [MLIR] Add `decomposeMixedStridesOrOffsets` and `decomposeMixedSizes`

[Frederik Gossen]

Author: Frederik Gossen
Date: 2022-06-23T19:08:36-04:00
New Revision: 02d29afd16d752eddb7949364343ea275fff6e30

URL: https://github.com/llvm/llvm-project/commit/02d29afd16d752eddb7949364343ea275fff6e30
DIFF: https://github.com/llvm/llvm-project/commit/02d29afd16d752eddb7949364343ea275fff6e30.diff

LOG: [MLIR] Add `decomposeMixedStridesOrOffsets` and `decomposeMixedSizes`

Add the reverse functions to the ViewLikeInterface's functions
`getMixedStrides`, `getMixedSizes`, and `getMixedOffsets`. The new functions
are useful to build view-like operations from an array of mixed static/dynamic
values.

Differential Revision: https://reviews.llvm.org/D128376

Added: 
    

Modified: 
    mlir/include/mlir/Interfaces/ViewLikeInterface.h
    mlir/lib/Interfaces/ViewLikeInterface.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.h b/mlir/include/mlir/Interfaces/ViewLikeInterface.h
index 8c18ab37017b1..18b3d4caba776 100644
--- a/mlir/include/mlir/Interfaces/ViewLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.h
@@ -49,6 +49,19 @@ SmallVector<OpFoldResult, 4> getMixedStrides(OffsetSizeAndStrideOpInterface op,
                                              ArrayAttr staticStrides,
                                              ValueRange strides);
 
+/// Decompose a vector of mixed static or dynamic strides/offsets into the
+/// corresponding pair of arrays. This is the inverse function of
+/// `getMixedStrides` and `getMixedOffsets`.
+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);
+
 namespace detail {
 LogicalResult verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op);
 

diff  --git a/mlir/lib/Interfaces/ViewLikeInterface.cpp b/mlir/lib/Interfaces/ViewLikeInterface.cpp
index 6f575cf87e969..4ecf5390edc75 100644
--- a/mlir/lib/Interfaces/ViewLikeInterface.cpp
+++ b/mlir/lib/Interfaces/ViewLikeInterface.cpp
@@ -242,3 +242,31 @@ mlir::getMixedStrides(OffsetSizeAndStrideOpInterface op,
   }
   return res;
 }
+
+static std::pair<ArrayAttr, SmallVector<Value>>
+decomposeMixedImpl(OpBuilder &b,
+                   const SmallVectorImpl<OpFoldResult> &mixedValues,
+                   const int64_t dynamicValuePlaceholder) {
+  SmallVector<int64_t> staticValues;
+  SmallVector<Value> dynamicValues;
+  for (const auto &it : mixedValues) {
+    if (it.is<Attribute>()) {
+      staticValues.push_back(it.get<Attribute>().cast<IntegerAttr>().getInt());
+    } else {
+      staticValues.push_back(ShapedType::kDynamicStrideOrOffset);
+      dynamicValues.push_back(it.get<Value>());
+    }
+  }
+  return {b.getI64ArrayAttr(staticValues), dynamicValues};
+}
+
+std::pair<ArrayAttr, SmallVector<Value>> mlir::decomposeMixedStridesOrOffsets(
+    OpBuilder &b, const SmallVectorImpl<OpFoldResult> &mixedValues) {
+  return decomposeMixedImpl(b, mixedValues, ShapedType::kDynamicStrideOrOffset);
+}
+
+std::pair<ArrayAttr, SmallVector<Value>>
+mlir::decomposeMixedSizes(OpBuilder &b,
+                          const SmallVectorImpl<OpFoldResult> &mixedValues) {
+  return decomposeMixedImpl(b, mixedValues, ShapedType::kDynamicSize);
+}