[Mlir-commits] [mlir] [mlir][vector] Update `castAwayContractionLeadingOneDim` to omit transposes solely on leading unit dims. (PR #85694)

Andrzej WarzyƄski llvmlistbot at llvm.org
Mon Mar 25 08:51:17 PDT 2024


================
@@ -399,13 +399,29 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
           transposeResults.push_back(targetExpr);
         }
       }
+
+      // Checks if only the outer, unit dimensions (of size 1) are permuted.
+      // Such transposes do not materially effect the underlying vector and can
+      // be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
+      bool tranposeNonOuterUnitDims = false;
+      for (auto [index, dim] :
+           llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
+        if (dim != static_cast<int64_t>(index) &&
+            operands[it.index()].getType().cast<ShapedType>().getDimSize(
+                index) != 1) {
+          tranposeNonOuterUnitDims = true;
+        }
+      }
----------------
banach-space wrote:

```suggestion
      auto operandShape = operands[it.index()].getType().cast<ShapedType>();
      for (auto [index, dim] :
           llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
        if (dim != static_cast<int64_t>(index) && operandShape.getDimSize(index) != 1) {
          tranposeNonOuterUnitDims = true;
          break;
        }
      }
```

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


More information about the Mlir-commits mailing list