[Mlir-commits] [mlir] [MLIR][Vector]Generalize DropUnitDimFromElementwiseOps (PR #92934)

Hugo Trachino llvmlistbot at llvm.org
Fri May 31 04:07:42 PDT 2024


================
@@ -1607,7 +1607,23 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
   }
 };
 
-/// For vectors with either leading or trailing unit dim, replaces:
+FailureOr<VectorType> dropNonScalableUnitDimType(VectorType inVecTy) {
+  int numUnitDimsDropped = 0;
+  auto inVecShape = inVecTy.getShape();
+  auto newVecBuilder = VectorType::Builder(inVecTy);
+  for (unsigned i = 0; i < inVecShape.size(); i++) {
+    if (inVecShape[i] == 1 && !inVecTy.getScalableDims()[i]) {
+      newVecBuilder.dropDim(i - numUnitDimsDropped);
+      numUnitDimsDropped++;
+    }
+  }
----------------
nujaa wrote:

Hi, thanks for the review. I pushed some change to handle scalable dims as well (that `dropDims` does under the hood). LMK if you are satisfied with the changes.

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


More information about the Mlir-commits mailing list