[mlir] Reuse the code between `getMixed*s()` funcs in ViewLikeInterface.cpp.

[jreiffers]

mlir/lib/Interfaces/ViewLikeInterface.cpp

183

Technically this is not the same anymore and this argument should be removed if it's safe to do so.

[Hardcode84]

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

123

staticValues[idx] == dynamicValueIndicator?

[pifon2a]

Address the comments.

[pifon2a]

mlir/lib/Interfaces/ViewLikeInterface.cpp

183

It can be done, but it needs some modifications for ViewLikeOpInterface.

[pifon2a]

mlir/lib/Interfaces/ViewLikeInterface.cpp

183

We cannot drop the op argument.

[pifon2a]

PTAL

[pifon2a]

Update.

[jreiffers]

mlir/include/mlir/Interfaces/ViewLikeInterface.td

240–263

Are the isDynamic... functions still needed?

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

136

It's probably better to just return this instead of the ArrayAttr.

1. It's easy to do the conversion at the call site
2. It lets you get rid of the builder argument.
3. SmallVector is easier to use and more type safe than ArrayAttr.

[pifon2a]

mlir/include/mlir/Interfaces/ViewLikeInterface.td

240–263

yes, they are used in a couple of places

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

136

While I agree with that, let's do that in a follow-up. I just moved the code here.

[frgossen]

mlir/include/mlir/Interfaces/ViewLikeInterface.td

248–263

Why do these have to be configurable? Is anyone using different placeholders?

[pifon2a]

mlir/include/mlir/Interfaces/ViewLikeInterface.td

248–263

No, they don't. This interface should ideally be refactored. Moreover, we can use one placeholder for dynamic sizes,strides and offsets, instead of using -1 for kDynamicSizes, but I guess there will be places where -1 was already hardcoded.

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

 
 /// Helper function to dispatch multiple OpFoldResults according to the behavior
 /// of `dispatchIndexOpFoldResult(OpFoldResult ofr` for a single OpFoldResult.
 void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
                                 SmallVectorImpl<Value> &dynamicVec,
                                 SmallVectorImpl<int64_t> &staticVec,
                                 int64_t sentinel);
 
+/// Return a vector of OpFoldResults given the special value
+/// that indicates whether of the value is dynamic or not.
+SmallVector<OpFoldResult, 4> getMixedValues(ArrayAttr staticValues,
+                                            ValueRange dynamicValues,
+                                            int64_t dynamicValueIndicator);
+
+/// Decompose a vector of mixed static or dynamic values into the corresponding
+/// pair of arrays. This is the inverse function of `getMixedValues`.
+std::pair<ArrayAttr, SmallVector<Value>>
+decomposeMixedValues(Builder &b,
+                     const SmallVectorImpl<OpFoldResult> &mixedValues,
+                     const int64_t dynamicValueIndicator);
+
 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr);
 
 /// Given a value, try to extract a constant Attribute. If this fails, return
 /// the original value.
 OpFoldResult getAsOpFoldResult(Value val);
 
 /// Given an array of values, try to extract a constant Attribute from each

mlir/include/mlir/Interfaces/ViewLikeInterface.td

InterfaceMethod<
       /*args=*/(ins "unsigned":$idx),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         ::llvm::APInt v = *(static_strides()
           .template getAsValueRange<::mlir::IntegerAttr>().begin() + idx);
         return ::mlir::ShapedType::isDynamicStrideOrOffset(v.getSExtValue());
       }]
     >,
-
+    StaticInterfaceMethod<
+      /*desc=*/"Return constant that indicates the offset is dynamic",
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getDynamicOffsetIndicator",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImpl=*/[{ return ::mlir::ShapedType::kDynamicStrideOrOffset; }]
+    >,
+    StaticInterfaceMethod<
+      /*desc=*/"Return constant that indicates the size is dynamic",
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getDynamicSizeIndicator",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImpl=*/[{ return ::mlir::ShapedType::kDynamicSize; }]
+    >,
+    StaticInterfaceMethod<
+      /*desc=*/"Return constant that indicates the stride is dynamic",
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getDynamicStrideIndicator",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImpl=*/[{ return ::mlir::ShapedType::kDynamicStrideOrOffset; }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Assert the offset `idx` is a static constant and return its value.
       }],
       /*retTy=*/"int64_t",
       /*methodName=*/"getStaticOffset",
       /*args=*/(ins "unsigned":$idx),
       /*methodBody=*/"",

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

 //===- StaticValueUtils.cpp - Utilities for dealing with static values ----===//
 //
 // 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/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/APSInt.h"
 
 namespace mlir {
 
 /// Helper function to dispatch an OpFoldResult into `staticVec` if:
 ///   a) it is an IntegerAttr

 /// no IndexAttr and that IndexType has no bitwidth.
 bool isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2) {
   auto cst1 = getConstantIntValue(ofr1), cst2 = getConstantIntValue(ofr2);
   if (cst1 && cst2 && *cst1 == *cst2)
     return true;
   auto v1 = ofr1.dyn_cast<Value>(), v2 = ofr2.dyn_cast<Value>();
   return v1 && v1 == v2;
 }
+
+/// Return a vector of OpFoldResults given the special value
+/// that indicates whether of the value is dynamic or not.
+SmallVector<OpFoldResult, 4> getMixedValues(ArrayAttr staticValues,
+                                            ValueRange dynamicValues,
+                                            int64_t dynamicValueIndicator) {
+  SmallVector<OpFoldResult, 4> res;
+  res.reserve(staticValues.size());
+  unsigned numDynamic = 0;
+  unsigned count = static_cast<unsigned>(staticValues.size());
+  for (unsigned idx = 0; idx < count; ++idx) {
+    APInt value = staticValues[idx].cast<IntegerAttr>().getValue();
+    res.push_back(value.getSExtValue() == dynamicValueIndicator
+                      ? OpFoldResult{dynamicValues[numDynamic++]}
+                      : OpFoldResult{staticValues[idx]});
+  }
+  return res;
+}
+
+std::pair<ArrayAttr, SmallVector<Value>>
+decomposeMixedValues(Builder &b,
+                     const SmallVectorImpl<OpFoldResult> &mixedValues,
+                     const int64_t dynamicValueIndicator) {
+  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(dynamicValueIndicator);
+      dynamicValues.push_back(it.get<Value>());
+    }
+  }
+  return {b.getI64ArrayAttr(staticValues), dynamicValues};
+}
+
 } // namespace mlir

mlir/lib/Interfaces/ViewLikeInterface.cpp

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;
+  return getMixedValues(staticOffsets, offsets, op.getDynamicOffsetIndicator());
 }
 
 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 getMixedValues(staticSizes, sizes, op.getDynamicSizeIndicator());
 }
 
 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;
-}
-
-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};
+  return getMixedValues(staticStrides, strides, op.getDynamicStrideIndicator());
 }
 
 std::pair<ArrayAttr, SmallVector<Value>> mlir::decomposeMixedStridesOrOffsets(
     OpBuilder &b, const SmallVectorImpl<OpFoldResult> &mixedValues) {
-  return decomposeMixedImpl(b, mixedValues, ShapedType::kDynamicStrideOrOffset);
+  return decomposeMixedValues(b, mixedValues,
+                              ShapedType::kDynamicStrideOrOffset);
 }
 
 std::pair<ArrayAttr, SmallVector<Value>>
 mlir::decomposeMixedSizes(OpBuilder &b,
                           const SmallVectorImpl<OpFoldResult> &mixedValues) {
-  return decomposeMixedImpl(b, mixedValues, ShapedType::kDynamicSize);
+  return decomposeMixedValues(b, mixedValues, ShapedType::kDynamicSize);
 }