diff --git a/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir b/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir deleted file mode 100644 index 739fb2fb8b68b..0000000000000 --- a/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir +++ /dev/null @@ -1,58 +0,0 @@ -// RUN: mlir-opt %s -split-input-file -test-bit-width-constrained-vector-linearize=target-vector-bitwidth=128 | FileCheck %s --check-prefixes=ALL,BW-128 -// RUN: mlir-opt %s -split-input-file -test-bit-width-constrained-vector-linearize=target-vector-bitwidth=0 | FileCheck %s --check-prefixes=ALL,BW-0 - -// A vector<2x2xf32> has inner-most dimension with 64-bits. Check that at -// bitwidth threshold 128 (>= 64), operations are linearized, and at -// bitwidth threshold 0 (< 64), operations are not linearized. - -// ALL-LABEL: test_result_bitwidth_64 -func.func @test_result_bitwidth_64(%arg0: vector<2x2xf32>) -> vector<2x2xf32> { - - // BW-128: arith.constant {{.*}} vector<4xf32> - // BW-0: arith.constant {{.*}} vector<2x2xf32> - %0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32> - - // BW-128: math.sin {{.*}} vector<4xf32> - // BW-0: math.sin {{.*}} vector<2x2xf32> - %1 = math.sin %arg0 : vector<2x2xf32> - - return %0 : vector<2x2xf32> -} - -// ----- - -// The size of the 'index' type is backend specific, so we cannot guarantee that -// the inner-most dimension below (of size 2*nbBits(index)) is below any bitwidth -// threshold. Test that operations with vectors of index type are not linearized. - -// ALL-LABEL: test_index_no_linearize -func.func @test_index_no_linearize(%arg0: vector<2x2xindex>, %arg1: vector<2x2xindex>) -> vector<2x2xindex> { - - // BW-128: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex> - // BW-0: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex> - %0 = arith.addi %arg0, %arg1 : vector<2x2xindex> - return %0 : vector<2x2xindex> -} - -// ----- - -// The logic for the insert op with regards to the bitwidth threshold is -// different to the other ops, so we test it here. Specifically, the logic -// is based on the bitwidth of the value to store. - -// ALL-LABEL: test_vector_insert -// ALL-SAME: (%[[DEST:.*]]: vector<2x8x4xf32>, %[[SRC:.*]]: vector<8x4xf32>) -> vector<2x8x4xf32> { -func.func @test_vector_insert(%arg0: vector<2x8x4xf32>, %arg1: vector<8x4xf32>) -> vector<2x8x4xf32> { - - // BW-128-DAG: %[[ARG_SRC:.*]] = vector.shape_cast %[[SRC]] : vector<8x4xf32> to vector<32xf32> - // BW-128-DAG: %[[ARG_DEST:.*]] = vector.shape_cast %[[DEST]] : vector<2x8x4xf32> to vector<64xf32> - // BW-128: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG_DEST]], %[[ARG_SRC]] - // BW-128: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<64xf32> to vector<2x8x4xf32> - // BW-128: return %[[RES]] : vector<2x8x4xf32> - - // BW-0: %[[RES:.*]] = vector.insert %[[SRC]], %[[DEST]] [0] : vector<8x4xf32> into vector<2x8x4xf32> - // BW-0: return %[[RES]] : vector<2x8x4xf32> - - %0 = vector.insert %arg1, %arg0[0]: vector<8x4xf32> into vector<2x8x4xf32> - return %0 : vector<2x8x4xf32> -} diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp index f4f32e9339870..03db9b7fa7a73 100644 --- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp +++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp @@ -837,126 +837,6 @@ struct TestVectorEmulateMaskedLoadStore final } }; -/// Get the set of operand/result types to check for sufficiently -/// small inner-most dimension size. -static SmallVector> -getTypeBitWidthBoundPairs(Operation *op, unsigned targetBitWidth) { - - if (auto insertOp = dyn_cast(op)) { - unsigned w = targetBitWidth < std::numeric_limits::max() - ? targetBitWidth + 1 - : targetBitWidth; - return {{insertOp.getValueToStoreType(), w}}; - } - - auto resultTypes = op->getResultTypes(); - SmallVector> resultsWithBitWidth; - resultsWithBitWidth.reserve(resultTypes.size()); - for (Type type : resultTypes) { - resultsWithBitWidth.push_back({type, targetBitWidth}); - } - return resultsWithBitWidth; -} - -/// If `type` is VectorType with trailing dimension of (bit) size greater than -/// or equal to `targetBitWidth`, its defining op is considered legal. -static bool -isNotLinearizableBecauseLargeInnerDimension(Type type, - unsigned targetBitWidth) { - - VectorType vecType = dyn_cast(type); - - // Not linearizable for reasons other than what this function checks. - if (!vecType || vecType.getRank() == 0) - return false; - - // The width of the type 'index' is unbounded (and therefore potentially above - // the target width). - if (vecType.getElementType().isIndex()) - return true; - - unsigned finalDimSize = vecType.getShape().back(); - unsigned nbBitsPerElm = vecType.getElementTypeBitWidth(); - unsigned trailingVecDimBitWidth = finalDimSize * nbBitsPerElm; - return trailingVecDimBitWidth >= targetBitWidth; -} - -static bool -isNotLinearizableBecauseLargeInnerDimension(Operation *op, - unsigned targetBitWidth) { - // Check on bitwidths. - SmallVector> toCheck = - getTypeBitWidthBoundPairs(op, targetBitWidth); - return llvm::any_of(toCheck, [&](std::pair typeWidth) { - return isNotLinearizableBecauseLargeInnerDimension(typeWidth.first, - typeWidth.second); - }); -} - -void populateWithBitWidthConstraints(TypeConverter &typeConverter, - ConversionTarget &target, - unsigned targetBitWidth) { - - // The general purpose definition of what ops are legal must come first. - populateForVectorLinearize(typeConverter, target); - - // Extend the set of legal ops to include those with large inner-most - // dimensions on selected operands/results. - target.markUnknownOpDynamicallyLegal( - [=](Operation *op) -> std::optional { - if (isNotLinearizableBecauseLargeInnerDimension(op, targetBitWidth)) { - return true; - } - return {}; - }); -} - -struct TestVectorBitWidthLinearize final - : public PassWrapper> { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBitWidthLinearize) - - TestVectorBitWidthLinearize() = default; - TestVectorBitWidthLinearize(const TestVectorBitWidthLinearize &pass) - : PassWrapper(pass) {} - - StringRef getArgument() const override { - return "test-bit-width-constrained-vector-linearize"; - } - StringRef getDescription() const override { - return "Linearizes ND vectors for N >= 2 into 1D vectors, with constraints " - "in inner-most dimension's bit width."; - } - void getDependentDialects(DialectRegistry ®istry) const override { - registry.insert(); - } - - Option targetVectorBitwidth{ - *this, "target-vector-bitwidth", - llvm::cl::desc( - "Minimum vector bitwidth to enable the flattening transformation"), - llvm::cl::init(std::numeric_limits::max())}; - void runOnOperation() override { - auto *context = &getContext(); - - TypeConverter typeConverter; - RewritePatternSet patterns(context); - ConversionTarget target(*context); - - populateWithBitWidthConstraints(typeConverter, target, - targetVectorBitwidth); - - vector::populateVectorLinearizeBasePatterns(typeConverter, target, - patterns); - - vector::populateVectorLinearizeShuffleLikeOpsPatterns(typeConverter, target, - patterns); - - if (failed(applyPartialConversion(getOperation(), target, - std::move(patterns)))) - return signalPassFailure(); - } -}; - struct TestVectorLinearize final : public PassWrapper> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorLinearize) @@ -1064,8 +944,6 @@ void registerTestVectorLowerings() { PassRegistration(); - PassRegistration(); - PassRegistration(); } } // namespace test