Add additional static value utils

Signed-off-by: Max Dawkins <[email protected]>

Author: Max191

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

@@ -92,6 +92,12 @@ getConstantIntValues(ArrayRef<OpFoldResult> ofrs);
 
 /// Return true if `ofr` is constant integer equal to `value`.
 bool isConstantIntValue(OpFoldResult ofr, int64_t value);
+/// Return true if all of `ofrs` are constant integers equal to `value`.
+bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value);
+/// Return true if all of `ofrs` are constant integers equal to the
+/// corresponding value in `values`.
+bool areConstantIntValues(ArrayRef<OpFoldResult> ofrs,
+                          ArrayRef<int64_t> values);
 
 /// Return true if ofr1 and ofr2 are the same integer constant attribute
 /// values or the same SSA value. Ignore integer bitwitdh and type mismatch

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -640,8 +640,6 @@ struct InsertOpInterface
 template <typename InsertOpTy>
 static bool insertSliceOpRequiresRead(InsertOpTy insertSliceOp,
                                       OpOperand &opOperand) {
-  RankedTensorType destType = insertSliceOp.getDestType();
-
   // The source is always read.
   if (opOperand == insertSliceOp.getSourceMutable())
     return true;
@@ -652,16 +650,12 @@ static bool insertSliceOpRequiresRead(InsertOpTy insertSliceOp,
   // Dest is not read if it is entirely overwritten. E.g.:
   // tensor.insert_slice %a into %t[0][10][1] : ... into tensor<10xf32>
   bool allOffsetsZero =
-      llvm::all_of(insertSliceOp.getMixedOffsets(),
-                   [](OpFoldResult ofr) { return isConstantIntValue(ofr, 0); });
-  bool sizesMatchDestSizes = llvm::all_of(
-      llvm::enumerate(insertSliceOp.getMixedSizes()), [&](const auto &it) {
-        return getConstantIntValue(it.value()) ==
-               destType.getDimSize(it.index());
-      });
+      llvm::all_of(insertSliceOp.getMixedOffsets(), isZeroIndex);
+  RankedTensorType destType = insertSliceOp.getDestType();
+  bool sizesMatchDestSizes =
+      areConstantIntValues(insertSliceOp.getMixedSizes(), destType.getShape());
   bool allStridesOne =
-      llvm::all_of(insertSliceOp.getMixedStrides(),
-                   [](OpFoldResult ofr) { return isConstantIntValue(ofr, 1); });
+      areAllConstantIntValue(insertSliceOp.getMixedStrides(), 1);
   return !(allOffsetsZero && sizesMatchDestSizes && allStridesOne);
 }
 

mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp

@@ -16,11 +16,6 @@ namespace mlir {
 namespace tensor {
 namespace {
 
-static bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value) {
-  return llvm::all_of(
-      ofrs, [&](OpFoldResult ofr) { return isConstantIntValue(ofr, value); });
-}
-
 /// Returns the number of shape sizes that is either dynamic or greater than 1.
 static int64_t getNumGtOneDims(ArrayRef<int64_t> shape) {
   return llvm::count_if(

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

@@ -10,6 +10,7 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/MathExtras.h"
 
 namespace mlir {
@@ -137,6 +138,22 @@ bool isConstantIntValue(OpFoldResult ofr, int64_t value) {
   return val && *val == value;
 }
 
+/// Return true if all of `ofrs` are constant integers equal to `value`.
+bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value) {
+  return llvm::all_of(
+      ofrs, [&](OpFoldResult ofr) { return isConstantIntValue(ofr, value); });
+}
+
+/// Return true if all of `ofrs` are constant integers equal to the
+/// corresponding value in `values`.
+bool areConstantIntValues(ArrayRef<OpFoldResult> ofrs,
+                          ArrayRef<int64_t> values) {
+  if (ofrs.size() != values.size())
+    return false;
+  std::optional<SmallVector<int64_t>> constOfrs = getConstantIntValues(ofrs);
+  return constOfrs && llvm::equal(constOfrs.value(), values);
+}
+
 /// Return true if ofr1 and ofr2 are the same integer constant attribute values
 /// or the same SSA value.
 /// Ignore integer bitwidth and type mismatch that come from the fact there is