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

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

@@ -1248,7 +1248,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.";
@@ -1443,7 +1443,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,
@@ -1495,7 +1495,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.";
@@ -1606,13 +1606,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)
     >,
@@ -169,15 +169,6 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
   ];
 
   let extraSharedClassDeclaration = [{
-    /// Return a vector of all in_bounds values as booleans (one per vector
-    /// transfer dimension).
-    ::llvm::SmallVector<bool> getInBoundsValues() {
-      ::llvm::SmallVector<bool> inBounds;
-      for (int64_t i = 0, e = $_op.getTransferRank(); i < e; ++i)
-        inBounds.push_back($_op.isDimInBounds(i));
-      return inBounds;
-    }
-
     /// Return the number of leading shaped dimensions (of the "source" operand)
     /// that do not participate in the permutation map.
     unsigned getLeadingShapedRank() {
@@ -241,7 +232,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

@@ -263,13 +263,6 @@ static void generateInBoundsCheck(
       });
 }
 
-/// Given an ArrayAttr, return a copy where the first element is dropped.
-static ArrayAttr dropFirstElem(OpBuilder &b, ArrayAttr attr) {
-  if (!attr)
-    return attr;
-  return ArrayAttr::get(b.getContext(), attr.getValue().drop_front());
-}
-
 /// Add the pass label to a vector transfer op if its rank is not the target
 /// rank.
 template <typename OpTy>
@@ -424,11 +417,11 @@ struct Strategy<TransferReadOp> {
     Location loc = xferOp.getLoc();
     auto bufferType = dyn_cast<ShapedType>(buffer.getType());
     auto vecType = dyn_cast<VectorType>(bufferType.getElementType());
-    auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
     auto newXferOp = b.create<vector::TransferReadOp>(
         loc, vecType, xferOp.getSource(), xferIndices,
         AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
-        xferOp.getPadding(), Value(), inBoundsAttr);
+        xferOp.getPadding(), Value(),
+        b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
 
     maybeApplyPassLabel(b, newXferOp, options.targetRank);
 
@@ -511,13 +504,12 @@ struct Strategy<TransferWriteOp> {
 
     Location loc = xferOp.getLoc();
     auto vec = b.create<memref::LoadOp>(loc, buffer, loadIndices);
-    auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
     auto source = loopState.empty() ? xferOp.getSource() : loopState[0];
     Type type = isTensorOp(xferOp) ? xferOp.getShapedType() : Type();
     auto newXferOp = b.create<vector::TransferWriteOp>(
         loc, type, vec, source, xferIndices,
         AffineMapAttr::get(unpackedPermutationMap(b, xferOp)), Value(),
-        inBoundsAttr);
+        b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
 
     maybeApplyPassLabel(b, newXferOp, options.targetRank);
 
@@ -1160,7 +1152,7 @@ struct ScalableTransposeTransferWriteConversion
               loopIterArgs.empty() ? writeOp.getSource() : loopIterArgs.front();
           auto newWriteOp = b.create<vector::TransferWriteOp>(
               loc, sliceVec, dest, xferIndices,
-              ArrayRef<bool>(writeOp.getInBoundsValues()).drop_front());
+              writeOp.getInBounds().drop_front());
           if (sliceMask)
             newWriteOp.getMaskMutable().assign(sliceMask);
 
@@ -1332,11 +1324,11 @@ struct UnrollTransferReadConversion
             getInsertionIndices(xferOp, insertionIndices);
             insertionIndices.push_back(rewriter.getIndexAttr(i));
 
-            auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
             auto newXferOp = b.create<vector::TransferReadOp>(
                 loc, newXferVecType, xferOp.getSource(), xferIndices,
                 AffineMapAttr::get(unpackedPermutationMap(b, xferOp)),
-                xferOp.getPadding(), Value(), inBoundsAttr);
+                xferOp.getPadding(), Value(),
+                b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
             maybeAssignMask(b, xferOp, newXferOp, i);
             return b.create<vector::InsertOp>(loc, newXferOp, vec,
                                               insertionIndices);
@@ -1467,7 +1459,6 @@ struct UnrollTransferWriteConversion
 
             auto extracted =
                 b.create<vector::ExtractOp>(loc, vec, extractionIndices);
-            auto inBoundsAttr = dropFirstElem(b, xferOp.getInBoundsAttr());
             Value xferVec;
             if (inputVectorTy.getRank() == 1) {
               // When target-rank=0, unrolling would causes the vector input
@@ -1481,7 +1472,7 @@ struct UnrollTransferWriteConversion
             auto newXferOp = b.create<vector::TransferWriteOp>(
                 loc, sourceType, xferVec, source, xferIndices,
                 AffineMapAttr::get(unpackedPermutationMap(b, xferOp)), Value(),
-                inBoundsAttr);
+                b.getDenseBoolArrayAttr(xferOp.getInBounds().drop_front()));
 
             maybeAssignMask(b, xferOp, newXferOp, i);
 

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());
@@ -902,7 +900,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

@@ -659,7 +659,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");
@@ -1399,7 +1399,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.
@@ -2432,7 +2432,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);
     });
@@ -2511,7 +2511,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();
@@ -2815,7 +2815,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)
@@ -2874,7 +2874,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);
@@ -3381,7 +3381,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);