[mlir][vector] Update the internal representation of in_bounds

This PR updates the internal representation of the `in_bounds` attribute
for `xfer_read`/`xfer_write` Ops. Currently we use `ArrayAttr` - that's
being updated to `DenseBoolArrayAttribute`.

Note that this means that the asm format of the `xfer_{read|_write}`
will change from:

```mlir
vector.transfer_read %arg0[%0, %1], %cst {in_bounds = [true], permutation_map = #map3} : memref<12x16xf32>, vector<8xf32>
```

to:
```mlir
vector.transfer_read %arg0[%0, %1], %cst {in_bounds = array<i1: true>, permutation_map = #map3} : memref<12x16xf32>, vector<8xf32>
```

Author: banach-space

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

@@ -1362,7 +1362,7 @@ def Vector_TransferReadOp :
                    AffineMapAttr:$permutation_map,
                    AnyType:$padding,
                    Optional<VectorOf<[I1]>>:$mask,
-                   BoolArrayAttr:$in_bounds)>,
+                   DenseBoolArrayAttr:$in_bounds)>,
     Results<(outs AnyVectorOfAnyRank:$vector)> {
 
   let summary = "Reads a supervector from memory into an SSA vector value.";
@@ -1557,7 +1557,7 @@ def Vector_TransferReadOp :
                    "Value":$source,
                    "ValueRange":$indices,
                    "AffineMapAttr":$permutationMapAttr,
-                   "ArrayAttr":$inBoundsAttr)>,
+                   "DenseBoolArrayAttr":$inBoundsAttr)>,
     /// 2. Builder that sets padding to zero and an empty mask (variant without attrs).
     OpBuilder<(ins "VectorType":$vectorType,
                    "Value":$source,
@@ -1609,7 +1609,7 @@ def Vector_TransferWriteOp :
                    Variadic<Index>:$indices,
                    AffineMapAttr:$permutation_map,
                    Optional<VectorOf<[I1]>>:$mask,
-                   BoolArrayAttr:$in_bounds)>,
+                   DenseBoolArrayAttr:$in_bounds)>,
     Results<(outs Optional<AnyRankedTensor>:$result)> {
 
   let summary = "The vector.transfer_write op writes a supervector to memory.";
@@ -1720,13 +1720,13 @@ def Vector_TransferWriteOp :
                    "ValueRange":$indices,
                    "AffineMapAttr":$permutationMapAttr,
                    "Value":$mask,
-                   "ArrayAttr":$inBoundsAttr)>,
+                   "DenseBoolArrayAttr":$inBoundsAttr)>,
     /// 2. Builder with type inference that sets an empty mask (variant with attrs).
     OpBuilder<(ins "Value":$vector,
                    "Value":$dest,
                    "ValueRange":$indices,
                    "AffineMapAttr":$permutationMapAttr,
-                   "ArrayAttr":$inBoundsAttr)>,
+                   "DenseBoolArrayAttr":$inBoundsAttr)>,
     /// 3. Builder with type inference that sets an empty mask (variant without attrs).
     OpBuilder<(ins "Value":$vector,
                    "Value":$dest,

mlir/include/mlir/Interfaces/VectorInterfaces.td

@@ -98,7 +98,7 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
         dimension whether it is in-bounds or not. (Broadcast dimensions are
         always in-bounds).
       }],
-      /*retTy=*/"::mlir::ArrayAttr",
+      /*retTy=*/"::mlir::ArrayRef<bool>",
       /*methodName=*/"getInBounds",
       /*args=*/(ins)
     >,
@@ -241,7 +241,7 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
       if ($_op.isBroadcastDim(dim))
         return true;
       auto inBounds = $_op.getInBounds();
-      return ::llvm::cast<::mlir::BoolAttr>(inBounds[dim]).getValue();
+      return inBounds[dim];
     }
 
     /// Helper function to account for the fact that `permutationMap` results

mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp

@@ -242,10 +242,11 @@ static void generateInBoundsCheck(
 }
 
 /// Given an ArrayAttr, return a copy where the first element is dropped.
-static ArrayAttr dropFirstElem(OpBuilder &b, ArrayAttr attr) {
+static DenseBoolArrayAttr dropFirstElem(OpBuilder &b, DenseBoolArrayAttr attr) {
   if (!attr)
     return attr;
-  return ArrayAttr::get(b.getContext(), attr.getValue().drop_front());
+  return DenseBoolArrayAttr::get(b.getContext(),
+                                 attr.asArrayRef().drop_front());
 }
 
 /// Add the pass label to a vector transfer op if its rank is not the target

mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp

@@ -497,8 +497,7 @@ struct LegalizeMultiTileTransferWriteAsStoreLoop
           loc, slice, writeOp.getSource(), ValueRange{storeRow, storeCol},
           AffineMapAttr::get(writeOp.getPermutationMap().dropResult(0)),
           sliceMask,
-          rewriter.getBoolArrayAttr(
-              ArrayRef<bool>(writeOp.getInBoundsValues()).drop_front()));
+          rewriter.getDenseBoolArrayAttr(writeOp.getInBounds().drop_front()));
     }
 
     rewriter.eraseOp(writeOp);
@@ -691,13 +690,12 @@ struct LiftIllegalVectorTransposeToMemory
         transposeOp.getPermutation(), getContext());
     auto transposedSubview = rewriter.create<memref::TransposeOp>(
         loc, readSubview, AffineMapAttr::get(transposeMap));
-    ArrayAttr inBoundsAttr = illegalRead.getInBoundsAttr();
+    DenseBoolArrayAttr inBoundsAttr = illegalRead.getInBoundsAttr();
     // - The `in_bounds` attribute
     if (inBoundsAttr) {
-      SmallVector<Attribute> inBoundsValues(inBoundsAttr.begin(),
-                                            inBoundsAttr.end());
+      SmallVector<bool> inBoundsValues(inBoundsAttr.asArrayRef());
       applyPermutationToVector(inBoundsValues, transposeOp.getPermutation());
-      inBoundsAttr = rewriter.getArrayAttr(inBoundsValues);
+      inBoundsAttr = rewriter.getDenseBoolArrayAttr(inBoundsValues);
     }
 
     VectorType legalReadType = resultType.clone(readType.getElementType());
@@ -990,7 +988,7 @@ struct LowerIllegalTransposeStoreViaZA
           rewriter.create<arith::AddIOp>(loc, transposedCol, writeIndices[1]);
       auto smeWrite = rewriter.create<vector::TransferWriteOp>(
           loc, tile, destTensorOrMemref, ValueRange{destRow, destCol},
-          transposeMap, subMask, writeOp.getInBounds());
+          transposeMap, subMask, writeOp.getInBoundsAttr());
 
       if (writeOp.hasPureTensorSemantics())
         destTensorOrMemref = smeWrite.getResult();

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -646,7 +646,7 @@ static Value buildVectorWrite(RewriterBase &rewriter, Value value,
   if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(write)) {
     auto maskedWriteOp = cast<vector::TransferWriteOp>(maskOp.getMaskableOp());
     SmallVector<bool> inBounds(maskedWriteOp.getVectorType().getRank(), true);
-    maskedWriteOp.setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
+    maskedWriteOp.setInBoundsAttr(rewriter.getDenseBoolArrayAttr(inBounds));
   }
 
   LDBG("vectorized op: " << *write << "\n");
@@ -1373,7 +1373,7 @@ vectorizeAsLinalgGeneric(RewriterBase &rewriter, VectorizationState &state,
     if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(read)) {
       SmallVector<bool> inBounds(readType.getRank(), true);
       cast<vector::TransferReadOp>(maskOp.getMaskableOp())
-          .setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
+          .setInBoundsAttr(rewriter.getDenseBoolArrayAttr(inBounds));
     }
 
     // 3.c. Not all ops support 0-d vectors, extract the scalar for now.
@@ -2406,7 +2406,7 @@ struct PadOpVectorizationWithTransferReadPattern
     rewriter.modifyOpInPlace(xferOp, [&]() {
       SmallVector<bool> inBounds(xferOp.getVectorType().getRank(), false);
       xferOp->setAttr(xferOp.getInBoundsAttrName(),
-                      rewriter.getBoolArrayAttr(inBounds));
+                      rewriter.getDenseBoolArrayAttr(inBounds));
       xferOp.getSourceMutable().assign(padOp.getSource());
       xferOp.getPaddingMutable().assign(padValue);
     });
@@ -2485,7 +2485,7 @@ struct PadOpVectorizationWithTransferWritePattern
     auto newXferOp = rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         xferOp, padOp.getSource().getType(), xferOp.getVector(),
         padOp.getSource(), xferOp.getIndices(), xferOp.getPermutationMapAttr(),
-        xferOp.getMask(), rewriter.getBoolArrayAttr(inBounds));
+        xferOp.getMask(), rewriter.getDenseBoolArrayAttr(inBounds));
     rewriter.replaceOp(trimPadding, newXferOp->getResult(0));
 
     return success();
@@ -2789,7 +2789,7 @@ LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
   Value res = rewriter.create<vector::TransferReadOp>(
       xferOp.getLoc(), vectorType, in, xferOp.getIndices(),
       xferOp.getPermutationMapAttr(), xferOp.getPadding(), xferOp.getMask(),
-      rewriter.getBoolArrayAttr(
+      rewriter.getDenseBoolArrayAttr(
           SmallVector<bool>(vectorType.getRank(), false)));
 
   if (maybeFillOp)
@@ -2848,7 +2848,7 @@ LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
   rewriter.create<vector::TransferWriteOp>(
       xferOp.getLoc(), vector, out, xferOp.getIndices(),
       xferOp.getPermutationMapAttr(), xferOp.getMask(),
-      rewriter.getBoolArrayAttr(
+      rewriter.getDenseBoolArrayAttr(
           SmallVector<bool>(vector.getType().getRank(), false)));
 
   rewriter.eraseOp(copyOp);
@@ -3348,7 +3348,7 @@ struct Conv1DGenerator
       SmallVector<bool> inBounds(maskShape.size(), true);
       auto xferOp = cast<VectorTransferOpInterface>(opToMask);
       xferOp->setAttr(xferOp.getInBoundsAttrName(),
-                      rewriter.getBoolArrayAttr(inBounds));
+                      rewriter.getDenseBoolArrayAttr(inBounds));
 
       SmallVector<OpFoldResult> mixedDims = vector::getMixedSizesXfer(
           cast<LinalgOp>(op).hasPureTensorSemantics(), opToMask, rewriter);

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -3792,7 +3792,7 @@ void ExtractStridedSliceOp::getCanonicalizationPatterns(
 void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            VectorType vectorType, Value source,
                            ValueRange indices, AffineMapAttr permutationMapAttr,
-                           /*optional*/ ArrayAttr inBoundsAttr) {
+                           /*optional*/ DenseBoolArrayAttr inBoundsAttr) {
   Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
   Value padding = builder.create<arith::ConstantOp>(
       result.location, elemType, builder.getZeroAttr(elemType));
@@ -3807,8 +3807,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
                            std::optional<ArrayRef<bool>> inBounds) {
   auto permutationMapAttr = AffineMapAttr::get(permutationMap);
   auto inBoundsAttr = (inBounds && !inBounds.value().empty())
-                          ? builder.getBoolArrayAttr(inBounds.value())
-                          : builder.getBoolArrayAttr(
+                          ? builder.getDenseBoolArrayAttr(inBounds.value())
+                          : builder.getDenseBoolArrayAttr(
                                 SmallVector<bool>(vectorType.getRank(), false));
   build(builder, result, vectorType, source, indices, permutationMapAttr,
         inBoundsAttr);
@@ -3823,8 +3823,8 @@ void TransferReadOp::build(OpBuilder &builder, OperationState &result,
       llvm::cast<ShapedType>(source.getType()), vectorType);
   auto permutationMapAttr = AffineMapAttr::get(permutationMap);
   auto inBoundsAttr = (inBounds && !inBounds.value().empty())
-                          ? builder.getBoolArrayAttr(inBounds.value())
-                          : builder.getBoolArrayAttr(
+                          ? builder.getDenseBoolArrayAttr(inBounds.value())
+                          : builder.getDenseBoolArrayAttr(
                                 SmallVector<bool>(vectorType.getRank(), false));
   build(builder, result, vectorType, source, indices, permutationMapAttr,
         padding,
@@ -3876,7 +3876,7 @@ static LogicalResult
 verifyTransferOp(VectorTransferOpInterface op, ShapedType shapedType,
                  VectorType vectorType, VectorType maskType,
                  VectorType inferredMaskType, AffineMap permutationMap,
-                 ArrayAttr inBounds) {
+                 ArrayRef<bool> inBounds) {
   if (op->hasAttr("masked")) {
     return op->emitOpError("masked attribute has been removed. "
                            "Use in_bounds instead.");
@@ -3949,8 +3949,7 @@ verifyTransferOp(VectorTransferOpInterface op, ShapedType shapedType,
            << AffineMapAttr::get(permutationMap)
            << " vs inBounds of size: " << inBounds.size();
   for (unsigned int i = 0, e = permutationMap.getNumResults(); i < e; ++i)
-    if (isa<AffineConstantExpr>(permutationMap.getResult(i)) &&
-        !llvm::cast<BoolAttr>(inBounds.getValue()[i]).getValue())
+    if (isa<AffineConstantExpr>(permutationMap.getResult(i)) && !inBounds[i])
       return op->emitOpError("requires broadcast dimensions to be in-bounds");
 
   return success();
@@ -4031,7 +4030,7 @@ ParseResult TransferReadOp::parse(OpAsmParser &parser, OperationState &result) {
   Attribute inBoundsAttr = result.attributes.get(inBoundsAttrName);
   if (!inBoundsAttr) {
     result.addAttribute(inBoundsAttrName,
-                        builder.getBoolArrayAttr(
+                        builder.getDenseBoolArrayAttr(
                             SmallVector<bool>(permMap.getNumResults(), false)));
   }
   if (parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
@@ -4159,7 +4158,7 @@ static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
     return failure();
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(op.getContext());
-  op.setInBoundsAttr(b.getBoolArrayAttr(newInBounds));
+  op.setInBoundsAttr(b.getDenseBoolArrayAttr(newInBounds));
   return success();
 }
 
@@ -4329,7 +4328,7 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value dest, ValueRange indices,
                             AffineMapAttr permutationMapAttr,
                             /*optional*/ Value mask,
-                            /*optional*/ ArrayAttr inBoundsAttr) {
+                            /*optional*/ DenseBoolArrayAttr inBoundsAttr) {
   Type resultType = llvm::dyn_cast<RankedTensorType>(dest.getType());
   build(builder, result, resultType, vector, dest, indices, permutationMapAttr,
         mask, inBoundsAttr);
@@ -4339,7 +4338,7 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value dest, ValueRange indices,
                             AffineMapAttr permutationMapAttr,
-                            /*optional*/ ArrayAttr inBoundsAttr) {
+                            /*optional*/ DenseBoolArrayAttr inBoundsAttr) {
   build(builder, result, vector, dest, indices, permutationMapAttr,
         /*mask=*/Value(), inBoundsAttr);
 }
@@ -4353,8 +4352,8 @@ void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
   auto permutationMapAttr = AffineMapAttr::get(permutationMap);
   auto inBoundsAttr =
       (inBounds && !inBounds.value().empty())
-          ? builder.getBoolArrayAttr(inBounds.value())
-          : builder.getBoolArrayAttr(SmallVector<bool>(
+          ? builder.getDenseBoolArrayAttr(inBounds.value())
+          : builder.getDenseBoolArrayAttr(SmallVector<bool>(
                 llvm::cast<VectorType>(vector.getType()).getRank(), false));
   build(builder, result, vector, dest, indices, permutationMapAttr,
         /*mask=*/Value(), inBoundsAttr);
@@ -4412,7 +4411,7 @@ ParseResult TransferWriteOp::parse(OpAsmParser &parser,
   Attribute inBoundsAttr = result.attributes.get(inBoundsAttrName);
   if (!inBoundsAttr) {
     result.addAttribute(inBoundsAttrName,
-                        builder.getBoolArrayAttr(
+                        builder.getDenseBoolArrayAttr(
                             SmallVector<bool>(permMap.getNumResults(), false)));
   }
   if (parser.resolveOperand(vectorInfo, vectorType, result.operands) ||
@@ -4765,7 +4764,7 @@ struct SwapExtractSliceOfTransferWrite
     auto newTransferWriteOp = rewriter.create<TransferWriteOp>(
         transferOp.getLoc(), transferOp.getVector(), newExtractOp.getResult(),
         transferOp.getIndices(), transferOp.getPermutationMapAttr(),
-        rewriter.getBoolArrayAttr(newInBounds));
+        rewriter.getDenseBoolArrayAttr(newInBounds));
     rewriter.modifyOpInPlace(insertOp, [&]() {
       insertOp.getSourceMutable().assign(newTransferWriteOp.getResult());
     });

mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp

@@ -22,15 +22,14 @@ using namespace mlir::vector;
 
 /// Transpose a vector transfer op's `in_bounds` attribute by applying reverse
 /// permutation based on the given indices.
-static ArrayAttr
-inverseTransposeInBoundsAttr(OpBuilder &builder, ArrayAttr attr,
+static DenseBoolArrayAttr
+inverseTransposeInBoundsAttr(OpBuilder &builder, ArrayRef<bool> inBounds,
                              const SmallVector<unsigned> &permutation) {
   SmallVector<bool> newInBoundsValues(permutation.size());
   size_t index = 0;
   for (unsigned pos : permutation)
-    newInBoundsValues[pos] =
-        cast<BoolAttr>(attr.getValue()[index++]).getValue();
-  return builder.getBoolArrayAttr(newInBoundsValues);
+    newInBoundsValues[pos] = inBounds[index++];
+  return builder.getDenseBoolArrayAttr(newInBoundsValues);
 }
 
 /// Extend the rank of a vector Value by `addedRanks` by adding outer unit
@@ -132,7 +131,7 @@ struct TransferReadPermutationLowering
     }
 
     // Transpose in_bounds attribute.
-    ArrayAttr newInBoundsAttr =
+    DenseBoolArrayAttr newInBoundsAttr =
         inverseTransposeInBoundsAttr(rewriter, op.getInBounds(), permutation);
 
     // Generate new transfer_read operation.
@@ -205,7 +204,7 @@ struct TransferWritePermutationLowering
                     });
 
     // Transpose in_bounds attribute.
-    ArrayAttr newInBoundsAttr =
+    DenseBoolArrayAttr newInBoundsAttr =
         inverseTransposeInBoundsAttr(rewriter, op.getInBounds(), permutation);
 
     // Generate new transfer_write operation.
@@ -298,7 +297,8 @@ struct TransferWriteNonPermutationLowering
     for (int64_t i = 0, e = op.getVectorType().getRank(); i < e; ++i) {
       newInBoundsValues.push_back(op.isDimInBounds(i));
     }
-    ArrayAttr newInBoundsAttr = rewriter.getBoolArrayAttr(newInBoundsValues);
+    DenseBoolArrayAttr newInBoundsAttr =
+        rewriter.getDenseBoolArrayAttr(newInBoundsValues);
     auto newWrite = rewriter.create<vector::TransferWriteOp>(
         op.getLoc(), newVec, op.getSource(), op.getIndices(),
         AffineMapAttr::get(newMap), newMask, newInBoundsAttr);
@@ -386,11 +386,8 @@ struct TransferOpReduceRank
 
     VectorType newReadType = VectorType::get(
         newShape, originalVecType.getElementType(), newScalableDims);
-    ArrayAttr newInBoundsAttr =
-        op.getInBounds()
-            ? rewriter.getArrayAttr(
-                  op.getInBoundsAttr().getValue().take_back(reducedShapeRank))
-            : ArrayAttr();
+    DenseBoolArrayAttr newInBoundsAttr = rewriter.getDenseBoolArrayAttr(
+        op.getInBounds().take_back(reducedShapeRank));
     Value newRead = rewriter.create<vector::TransferReadOp>(
         op.getLoc(), newReadType, op.getSource(), op.getIndices(),
         AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

@@ -248,10 +248,9 @@ struct CastAwayTransferReadLeadingOneDim
         AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(), newResults,
                        rewriter.getContext());
 
-    ArrayAttr inBoundsAttr;
-    if (read.getInBounds())
-      inBoundsAttr = rewriter.getArrayAttr(
-          read.getInBoundsAttr().getValue().take_back(newType.getRank()));
+    DenseBoolArrayAttr inBoundsAttr;
+    inBoundsAttr = rewriter.getDenseBoolArrayAttr(
+        read.getInBoundsAttr().asArrayRef().take_back(newType.getRank()));
 
     Value mask = Value();
     if (read.getMask()) {
@@ -302,10 +301,9 @@ struct CastAwayTransferWriteLeadingOneDim
         AffineMap::get(oldMap.getNumDims(), oldMap.getNumSymbols(), newResults,
                        rewriter.getContext());
 
-    ArrayAttr inBoundsAttr;
-    if (write.getInBounds())
-      inBoundsAttr = rewriter.getArrayAttr(
-          write.getInBoundsAttr().getValue().take_back(newType.getRank()));
+    DenseBoolArrayAttr inBoundsAttr;
+    inBoundsAttr = rewriter.getDenseBoolArrayAttr(
+        write.getInBoundsAttr().asArrayRef().take_back(newType.getRank()));
 
     auto newVector = rewriter.create<vector::ExtractOp>(
         write.getLoc(), write.getVector(), splatZero(dropDim));