Replace iree_compiler::VscaleRange with vector::VscaleRange (NFC)

Signed-off-by: Benjamin Maxwell <benjamin.maxwell@arm.com>

compiler/src/iree/compiler/Codegen/Common/FoldTensorSubsetIntoVectorTransferOps.cpp

@@ -140,7 +140,7 @@ static bool isDestinationFullyOverwritten(vector::TransferWriteOp writeOp) {
   if (writeOp.getMask())
     return false;
 
-  std::optional<iree_compiler::VscaleRange> vscaleRange;
+  std::optional<vector::VscaleRange> vscaleRange;
   auto vecType = writeOp.getVectorType();
   if (vecType.isScalable()) {
     auto targetAttr =

compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp

@@ -53,7 +53,7 @@ inferSizesFromIR(linalg::LinalgOp linalgOp, std::optional<OpResult> opResult) {
     VEC_DBGS() << '\n';
   });
 
-  std::optional<VscaleRange> vscaleRange;
+  std::optional<vector::VscaleRange> vscaleRange;
   if (!opResult) {
     // Note: Inferring scalable sizes is not supported is `opResult` is set
     // (which is used to compute sizes for tensor.pack/unpack).
@@ -325,14 +325,6 @@ class GenericVectorizationPass final
   void runOnOperation() override;
 };
 
-/// Converts from iree_compiler::VscaleRange to vector::VscaleRange.
-static std::optional<vector::VscaleRange>
-toVectorVscaleRange(std::optional<iree_compiler::VscaleRange> vscaleRange) {
-  if (!vscaleRange.has_value())
-    return std::nullopt;
-  return vector::VscaleRange{vscaleRange->min, vscaleRange->max};
-}
-
 void GenericVectorizationPass::runOnOperation() {
   MLIRContext *context = &getContext();
   auto funcOp = getOperation();
@@ -392,8 +384,7 @@ void GenericVectorizationPass::runOnOperation() {
     auto targetAttr =
         iree_compiler::IREE::HAL::ExecutableTargetAttr::lookup(funcOp);
     auto vscaleRange = iree_compiler::getDefaultVscaleRange(targetAttr);
-    vector::eliminateVectorMasks(rewriter, funcOp,
-                                 toVectorVscaleRange(vscaleRange));
+    vector::eliminateVectorMasks(rewriter, funcOp, vscaleRange);
   }
 
   {

compiler/src/iree/compiler/Codegen/Common/HoistStaticallyBoundAllocations.cpp

@@ -36,7 +36,7 @@ void HoistStaticallyBoundAllocationsPass::runOnOperation() {
   auto funcOp = getOperation();
   IRRewriter rewriter(funcOp->getContext());
 
-  std::optional<VscaleRange> vscaleRange;
+  std::optional<vector::VscaleRange> vscaleRange;
   if (this->vscaleMax != 0 && this->vscaleMin <= this->vscaleMax)
     vscaleRange = {this->vscaleMin, this->vscaleMax};
 

compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp

@@ -111,7 +111,8 @@ template <typename AllocLikeOpType>
 std::optional<Value> hoistOneStaticallyBoundAllocation(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder, Location loc,
     MemRefType allocLikeType, ValueRange dynamicSizes,
-    std::optional<uint64_t> alignment, std::optional<VscaleRange> vscaleRange) {
+    std::optional<uint64_t> alignment,
+    std::optional<vector::VscaleRange> vscaleRange) {
   IntegerAttr alignmentAttr =
       alignment ? builder.getI64IntegerAttr(alignment.value()) : nullptr;
   // For static case just create a new allocation in the entry block of the same
@@ -142,8 +143,8 @@ std::optional<Value> hoistOneStaticallyBoundAllocation(
       // Scalability supported: Allocations could be scalable.
       FailureOr<vector::ConstantOrScalableBound> ub =
           vector::ScalableValueBoundsConstraintSet::computeScalableBound(
-              value, std::nullopt, vscaleRange->min, vscaleRange->max,
-              presburger::BoundType::UB);
+              value, std::nullopt, vscaleRange->vscaleMin,
+              vscaleRange->vscaleMax, presburger::BoundType::UB);
       if (failed(ub))
         return failure();
 
@@ -234,7 +235,8 @@ std::optional<Value> hoistOneStaticallyBoundAllocation(
 template <typename AllocLikeOpType>
 std::optional<Value> hoistOneStaticallyBoundAllocation(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder,
-    AllocLikeOpType allocLikeOp, std::optional<VscaleRange> vscaleRange) {
+    AllocLikeOpType allocLikeOp,
+    std::optional<vector::VscaleRange> vscaleRange) {
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPoint(allocLikeOp);
   return hoistOneStaticallyBoundAllocation<AllocLikeOpType>(
@@ -254,7 +256,7 @@ static bool isUseReplaceableWithSubview(OpOperand &use) {
 template <typename AllocLikeOpType>
 void hoistStaticallyBoundAllocationsInFunc(
     RewriterBase &rewriter, mlir::FunctionOpInterface funcOp,
-    std::optional<VscaleRange> vscaleRange) {
+    std::optional<vector::VscaleRange> vscaleRange) {
   SmallVector<AllocLikeOpType> allocLikeOps;
 
   // Collect all allocLikes that are hoistable.
@@ -312,26 +314,30 @@ template std::optional<Value>
 hoistOneStaticallyBoundAllocation<memref::AllocOp>(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder, Location loc,
     MemRefType allocLikeType, ValueRange dynamicSizes,
-    std::optional<uint64_t> alignment, std::optional<VscaleRange> vscaleRange);
+    std::optional<uint64_t> alignment,
+    std::optional<vector::VscaleRange> vscaleRange);
 template std::optional<Value>
 hoistOneStaticallyBoundAllocation<memref::AllocaOp>(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder, Location loc,
     MemRefType allocLikeType, ValueRange dynamicSizes,
-    std::optional<uint64_t> alignment, std::optional<VscaleRange> vscaleRange);
+    std::optional<uint64_t> alignment,
+    std::optional<vector::VscaleRange> vscaleRange);
 template std::optional<Value>
 hoistOneStaticallyBoundAllocation<memref::AllocOp>(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder,
-    memref::AllocOp allocLikeOp, std::optional<VscaleRange> vscaleRange);
+    memref::AllocOp allocLikeOp,
+    std::optional<vector::VscaleRange> vscaleRange);
 template std::optional<Value>
 hoistOneStaticallyBoundAllocation<memref::AllocaOp>(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder,
-    memref::AllocaOp allocLikeOp, std::optional<VscaleRange> vscaleRange);
+    memref::AllocaOp allocLikeOp,
+    std::optional<vector::VscaleRange> vscaleRange);
 template void hoistStaticallyBoundAllocationsInFunc<memref::AllocOp>(
     RewriterBase &rewriter, mlir::FunctionOpInterface funcOp,
-    std::optional<VscaleRange> vscaleRange);
+    std::optional<vector::VscaleRange> vscaleRange);
 template void hoistStaticallyBoundAllocationsInFunc<memref::AllocaOp>(
     RewriterBase &rewriter, mlir::FunctionOpInterface funcOp,
-    std::optional<VscaleRange> vscaleRange);
+    std::optional<vector::VscaleRange> vscaleRange);
 
 //===---------------------------------------------------------------------===//
 // Lowering `flow.dispatch.workgroup_count_from_slice` operation.

compiler/src/iree/compiler/Codegen/Transforms/Transforms.h

@@ -53,7 +53,7 @@ std::optional<Value> hoistOneStaticallyBoundAllocation(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder, Location loc,
     MemRefType allocaType, ValueRange dynamicSizes,
     std::optional<uint64_t> alignment,
-    std::optional<VscaleRange> vscaleRange = std::nullopt);
+    std::optional<vector::VscaleRange> vscaleRange = std::nullopt);
 
 /// Hoists `allocaOp` to the entry block of the function if the size is
 /// statically bounded. For a static allocation, it returns an allocation
@@ -65,14 +65,14 @@ template <typename AllocLikeOpType>
 std::optional<Value> hoistOneStaticallyBoundAllocation(
     mlir::FunctionOpInterface funcOp, OpBuilder &builder,
     AllocLikeOpType allocaOp,
-    std::optional<VscaleRange> vscaleRange = std::nullopt);
+    std::optional<vector::VscaleRange> vscaleRange = std::nullopt);
 
 /// Traverse funcOp and try to hoist every AllocaOp to the entry block of the
 /// function if the size is statically bounded.
 template <typename AllocLikeOpType>
 void hoistStaticallyBoundAllocationsInFunc(
     RewriterBase &rewriter, mlir::FunctionOpInterface funcOp,
-    std::optional<VscaleRange> vscaleRange = std::nullopt);
+    std::optional<vector::VscaleRange> vscaleRange = std::nullopt);
 
 /// Insert patterns to perform folding of AffineMinOp by matching the
 /// pattern generated by tile and distribute. Try to fold a affine.min op by

compiler/src/iree/compiler/Codegen/Utils/BUILD.bazel

@@ -66,6 +66,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:TilingInterface",
         "@llvm-project//mlir:TransformUtils",
         "@llvm-project//mlir:VectorDialect",
+        "@llvm-project//mlir:VectorTransforms",
         "@llvm-project//mlir:ViewLikeInterface",
     ],
 )

compiler/src/iree/compiler/Codegen/Utils/CMakeLists.txt

@@ -50,6 +50,7 @@ iree_cc_library(
     MLIRTilingInterface
     MLIRTransformUtils
     MLIRVectorDialect
+    MLIRVectorTransforms
     MLIRViewLikeInterface
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect

compiler/src/iree/compiler/Codegen/Utils/Utils.cpp

@@ -1167,21 +1167,21 @@ bool hasFusedLeadingOp(linalg::LinalgOp rootOp) {
   return llvm::any_of(backwardSlice, llvm::IsaPred<linalg::LinalgOp>);
 }
 
-std::optional<VscaleRange>
+std::optional<vector::VscaleRange>
 getDefaultVscaleRange(IREE::HAL::ExecutableTargetAttr targetAttr) {
   if (isAArch64(targetAttr)) {
     // On AArch64 the scalable vector length will always be between 128-bit and
     // 2048-bit. This works out as a vscale range of 1 to 16. See:
     // https://developer.arm.com/Architectures/Scalable%20Vector%20Extensions
-    return VscaleRange{1, 16};
+    return vector::VscaleRange{1, 16};
   }
   // TODO: Implement for other architectures.
   return std::nullopt;
 }
 
 FailureOr<DimBoundSize>
 computeDimUpperBound(Value shapedValue, unsigned dimNum,
-                     std::optional<VscaleRange> vscaleRange,
+                     std::optional<vector::VscaleRange> vscaleRange,
                      RoundUpVscaleMultiple roundUp) {
   if (!vscaleRange.has_value()) {
     FailureOr<int64_t> maybeDimBoundSize =
@@ -1196,8 +1196,8 @@ computeDimUpperBound(Value shapedValue, unsigned dimNum,
   FailureOr<DimBound> maybeDimBound =
       vector::ScalableValueBoundsConstraintSet::computeScalableBound(
           shapedValue, dimNum,
-          /*vscaleMin=*/vscaleRange->min,
-          /*vscaleMax=*/vscaleRange->max, presburger::BoundType::UB);
+          /*vscaleMin=*/vscaleRange->vscaleMin,
+          /*vscaleMax=*/vscaleRange->vscaleMax, presburger::BoundType::UB);
   if (failed(maybeDimBound))
     return failure();
   auto boundSize = maybeDimBound->getSize();

compiler/src/iree/compiler/Codegen/Utils/Utils.h

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
 #include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
@@ -226,12 +227,7 @@ void sinkOpsInCFG(const SmallVector<Operation *> &allocs,
 // the inputs.
 bool hasFusedLeadingOp(linalg::LinalgOp rootOp);
 
-struct VscaleRange {
-  unsigned min;
-  unsigned max;
-};
-
-std::optional<VscaleRange>
+std::optional<vector::VscaleRange>
 getDefaultVscaleRange(IREE::HAL::ExecutableTargetAttr targetAttr);
 
 using DimBound = vector::ConstantOrScalableBound;
@@ -246,7 +242,7 @@ enum class RoundUpVscaleMultiple { No, Yes };
 /// bound can be a scalable quantity.
 FailureOr<DimBoundSize>
 computeDimUpperBound(Value shapedValue, unsigned dimNum,
-                     std::optional<VscaleRange> vscaleRange,
+                     std::optional<vector::VscaleRange> vscaleRange,
                      RoundUpVscaleMultiple = RoundUpVscaleMultiple::No);
 
 } // namespace mlir::iree_compiler