[Mlir-commits] [mlir] [mlir][Linalg] Add dropI64ArrayAttrElem helper to reduce duplication (NFC) (PR #174279)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: Nick Kreeger (nkreeger)

<details>
<summary>Changes</summary>

Adds a utility function to drop an element from a DenseIntElementsAttr at a specified index. This consolidates duplicate code in Conv2D rank-reduction transformations.

---
Full diff: https://github.com/llvm/llvm-project/pull/174279.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Linalg/Utils/Utils.h (+7) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp (+8-19) 
- (modified) mlir/lib/Dialect/Linalg/Utils/Utils.cpp (+8) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
index 9da01f30b52d2..ae3d2b2e5ddfa 100644
--- a/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
@@ -392,6 +392,13 @@ SmallVector<NamedAttribute> getPrunedAttributeList(OpTy op) {
   return getPrunedAttributeList(op, elidedAttrs);
 }
 
+/// Creates a new I64 array attribute by dropping the element at the specified
+/// index from the input array attribute. This is useful for rank reduction
+/// operations where dimensions need to be removed from strides, dilations, etc.
+DenseIntElementsAttr dropI64ArrayAttrElem(OpBuilder &builder,
+                                          DenseIntElementsAttr attr,
+                                          unsigned index);
+
 } // namespace linalg
 } // namespace mlir
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 96cc378f6c21a..37d6f6d079799 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1483,16 +1483,10 @@ FailureOr<Conv1DOp> DownscaleSizeOneWindowed2DConvolution<Conv2DOp, Conv1DOp>::
       rewriter, loc, output, newOutputType);
 
   // Rank-reduce strides and dilations too.
-  // TODO: dropDim 1-liner helper.
-  auto strides =
-      llvm::to_vector<4>(convOp.getStrides().template getValues<int64_t>());
-  strides.erase(strides.begin() + (removeH ? 0 : 1));
-  auto stridesAttr = rewriter.getI64VectorAttr(strides);
-
-  auto dilations =
-      llvm::to_vector<4>(convOp.getDilations().template getValues<int64_t>());
-  dilations.erase(dilations.begin() + (removeH ? 0 : 1));
-  auto dilationsAttr = rewriter.getI64VectorAttr(dilations);
+  auto stridesAttr =
+      dropI64ArrayAttrElem(rewriter, convOp.getStrides(), removeH ? 0 : 1);
+  auto dilationsAttr =
+      dropI64ArrayAttrElem(rewriter, convOp.getDilations(), removeH ? 0 : 1);
 
   auto conv1DOp = Conv1DOp::create(
       rewriter, loc, newOutputType, ValueRange{newInput, newKernel},
@@ -1571,15 +1565,10 @@ DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite(
       rewriter, loc, output, newOutputType);
 
   // Rank-reduce strides and dilations too.
-  // TODO: dropDim 1-liner helper.
-  auto strides = llvm::to_vector<4>(convOp.getStrides().getValues<int64_t>());
-  strides.erase(strides.begin() + (removeH ? 0 : 1));
-  auto stridesAttr = rewriter.getI64VectorAttr(strides);
-
-  auto dilations =
-      llvm::to_vector<4>(convOp.getDilations().getValues<int64_t>());
-  dilations.erase(dilations.begin() + (removeH ? 0 : 1));
-  auto dilationsAttr = rewriter.getI64VectorAttr(dilations);
+  auto stridesAttr =
+      dropI64ArrayAttrElem(rewriter, convOp.getStrides(), removeH ? 0 : 1);
+  auto dilationsAttr =
+      dropI64ArrayAttrElem(rewriter, convOp.getDilations(), removeH ? 0 : 1);
 
   auto conv1DOp = DepthwiseConv1DNwcWcOp::create(
       rewriter, loc, newOutputType, ValueRange{newInput, newKernel},
diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 2718124251c18..a64b82c183c1d 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -2573,5 +2573,13 @@ getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes) {
   return reassociation;
 }
 
+DenseIntElementsAttr dropI64ArrayAttrElem(OpBuilder &builder,
+                                          DenseIntElementsAttr attr,
+                                          unsigned index) {
+  auto values = llvm::to_vector<4>(attr.getValues<int64_t>());
+  values.erase(values.begin() + index);
+  return builder.getI64VectorAttr(values);
+}
+
 } // namespace linalg
 } // namespace mlir

``````````

</details>


https://github.com/llvm/llvm-project/pull/174279