[Mlir-commits] [mlir] [mlir][vector] extend `createReadOrMaskedRead`/`createWriteOrMaskedWrite` with permutation map support (PR #202766)

[Andrzej Warzyński]

================
@@ -434,30 +457,47 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
   int64_t vecToReadRank = vecToReadTy.getRank();
   auto vecToReadShape = vecToReadTy.getShape();
 
-  assert(sourceShape.size() == static_cast<size_t>(vecToReadRank) &&
-         "expected same ranks.");
+  size_t expectedSourceRank =
+      permutationMap ? permutationMap.getNumDims() : vecToReadRank;
+  assert(sourceShape.size() == expectedSourceRank &&
+         "expected source rank to match permutation map dims or vector rank.");
   assert((!padValue.has_value() ||
           padValue.value().getType() == sourceShapedType.getElementType()) &&
          "expected same pad element type to match source element type");
 
-  auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
   SmallVector<bool> inBoundsVal(vecToReadRank, true);
 
   if (useInBoundsInsteadOfMasking) {
-    // Update the inBounds attribute.
-    // FIXME: This computation is too weak - it ignores the read indices.
-    for (unsigned i = 0; i < vecToReadRank; i++)
-      inBoundsVal[i] = (sourceShape[i] == vecToReadShape[i]) &&
-                       ShapedType::isStatic(sourceShape[i]);
+    if (permutationMap) {
+      // FIXME: This computation is too weak - it ignores the read indices.
+      inBoundsVal = computeInBoundsFromPermutationMap(
+          permutationMap, vecToReadTy, cast<MemRefType>(source.getType()));
+    } else {
+      // Update the inBounds attribute.
----------------
banach-space wrote:

This comment also applies to both cases, right?

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