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

[Andrzej Warzyński]

================
@@ -434,29 +457,54 @@ 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.");
+  // The permutation map maps the source's index space to the vector's, so its
+  // dims must match the source rank and its results the vector rank. Without a
+  // map, a minor identity is implied, requiring the two ranks to match.
+  assert(sourceShape.size() == (permutationMap
+                                    ? permutationMap.getNumDims()
+                                    : static_cast<size_t>(vecToReadRank)) &&
+         "expected source rank to match permutation map dims or vector rank.");
+  assert((!permutationMap || permutationMap.getNumResults() ==
+                                 static_cast<size_t>(vecToReadRank)) &&
+         "expected permutation map results to match 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) {
+      // Update the inBounds attribute.
+      // FIXME: This computation is too weak - it ignores the read indices.
+      inBoundsVal = computeInBoundsFromPermutationMap(
+          permutationMap, vecToReadTy, cast<ShapedType>(source.getType()));
+    } else {
+      // 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]);
+    }
+  }
+  // The transfer op expects one index per source dimension.
+  assert(
+      (customIndices.empty() || customIndices.size() == sourceShape.size()) &&
+      "expected as many custom indices as source dims.");
+  SmallVector<Value> indices;
+  if (customIndices.empty()) {
+    auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
+    indices.assign(sourceShape.size(), zero);
+  } else {
+    indices.assign(customIndices.begin(), customIndices.end());
----------------
banach-space wrote:

[nit] I personally find the ternary operator much clearer for situations like this.

```suggestion
    customIndices.empty()
        ? indices.assign(sourceShape.size(),
                         arith::ConstantIndexOp::create(builder, loc, 0))
        : indices.assign(customIndices.begin(), customIndices.end());
```

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