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

[Andrzej Warzyński]

================
@@ -434,30 +457,45 @@ Value vector::createReadOrMaskedRead(OpBuilder &builder, Location loc,
   int64_t vecToReadRank = vecToReadTy.getRank();
   auto vecToReadShape = vecToReadTy.getShape();
 
-  assert(sourceShape.size() == static_cast<size_t>(vecToReadRank) &&
+  assert((permutationMap ||
+          sourceShape.size() == static_cast<size_t>(vecToReadRank)) &&
          "expected same ranks.");
   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) {
+      inBoundsVal = computeInBoundsFromPermutationMap(
+          permutationMap, vecToReadTy, cast<MemRefType>(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]);
+    }
+  }
+  SmallVector<Value> indices;
+  if (customIndices.empty()) {
+    auto zero = arith::ConstantIndexOp::create(builder, loc, 0);
+    indices.assign(vecToReadRank, zero);
+  } else {
+    indices.assign(customIndices.begin(), customIndices.end());
   }
-  SmallVector<Value> indices(vecToReadRank, zero);
   auto transferReadOp =
-      vector::TransferReadOp::create(builder, loc,
-                                     /*vectorType=*/vecToReadTy,
-                                     /*source=*/source,
-                                     /*indices=*/indices,
-                                     /*padding=*/padValue,
-                                     /*inBounds=*/inBoundsVal);
+      permutationMap
+          ? vector::TransferReadOp::create(builder, loc, vecToReadTy, source,
+                                           indices, padValue, permutationMap,
+                                           inBoundsVal)
+          : vector::TransferReadOp::create(builder, loc,
+                                           /*vectorType=*/vecToReadTy,
+                                           /*source=*/source,
+                                           /*indices=*/indices,
+                                           /*padding=*/padValue,
+                                           /*inBounds=*/inBoundsVal);
----------------
banach-space wrote:

I would just add a new one. That said, we should probably check whether there's room for some clean-up. We probably don't require all those builders.


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