Skip to content

Commit

Permalink
[mlir][ArmSME] Move vector.extract/insert lowerings to vector-to-arm-…
Browse files Browse the repository at this point in the history 
…sme (NFC) (llvm#72852)

These were placed in LegalizeForLLVMExport.cpp, which is the wrong stage
for these, as these lower to high-level ArmSME ops, not intrinsics.
  • Loading branch information
@MacDue
MacDue authored Nov 20, 2023
1 parent 28b5054 commit c4c52d4
Show file tree
Hide file tree
Showing 2 changed files with 114 additions and 115 deletions.
114 changes: 113 additions & 1 deletion mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -573,6 +573,116 @@ struct VectorOuterProductToArmSMELowering
}
};

/// Lower `vector.extract` using `arm_sme.move_tile_slice_to_vector`.
///
/// Example:
/// ```
/// %el = vector.extract %tile[%row, %col]: i32 from vector<[4]x[4]xi32>
/// ```
/// Becomes:
/// ```
/// %slice = arm_sme.move_tile_slice_to_vector %tile[%row]
/// : vector<[4]xi32> from vector<[4]x[4]xi32>
/// %el = vector.extract %slice[%col] : i32 from vector<[4]xi32>
/// ```
struct VectorExtractToArmSMELowering
: public OpRewritePattern<vector::ExtractOp> {
using OpRewritePattern<vector::ExtractOp>::OpRewritePattern;

LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
PatternRewriter &rewriter) const override {
VectorType sourceType = extractOp.getSourceVectorType();
if (!arm_sme::isValidSMETileVectorType(sourceType))
return failure();

auto loc = extractOp.getLoc();
auto position = extractOp.getMixedPosition();

Value sourceVector = extractOp.getVector();

// Extract entire vector. Should be handled by folder, but just to be safe.
if (position.empty()) {
rewriter.replaceOp(extractOp, sourceVector);
return success();
}

Value sliceIndex = vector::getAsValues(rewriter, loc, position[0]).front();
auto moveTileSliceToVector =
rewriter.create<arm_sme::MoveTileSliceToVectorOp>(loc, sourceVector,
sliceIndex);

if (position.size() == 1) {
// Single index case: Extracts a 1D slice.
rewriter.replaceOp(extractOp, moveTileSliceToVector);
return success();
}

// Two indices case: Extracts a single element.
assert(position.size() == 2);
rewriter.replaceOpWithNewOp<vector::ExtractOp>(
extractOp, moveTileSliceToVector, position[1]);

return success();
}
};

/// Lower `vector.insert` using `arm_sme.move_vector_to_tile_slice` and
/// `arm_sme.move_tile_slice_to_vector`.
///
/// Example:
/// ```
/// %new_tile = vector.insert %el, %tile[%row, %col]
/// : i32 into vector<[4]x[4]xi32>
/// ```
/// Becomes:
/// ```
/// %slice = arm_sme.move_tile_slice_to_vector %tile[%row]
/// : vector<[4]xi32> from vector<[4]x[4]xi32>
/// %new_slice = vector.insert %el, %slice[%col] : i32 into vector<[4]xi32>
/// %new_tile = arm_sme.move_vector_to_tile_slice %new_slice, %tile, %row
/// : vector<[4]xi32> into vector<[4]x[4]xi32>
/// ```
struct VectorInsertToArmSMELowering
: public OpRewritePattern<vector::InsertOp> {
using OpRewritePattern<vector::InsertOp>::OpRewritePattern;

LogicalResult matchAndRewrite(vector::InsertOp insertOp,
PatternRewriter &rewriter) const override {
VectorType resultType = insertOp.getResult().getType();

if (!arm_sme::isValidSMETileVectorType(resultType))
return failure();

auto loc = insertOp.getLoc();
auto position = insertOp.getMixedPosition();

Value source = insertOp.getSource();

// Overwrite entire vector with value. Should be handled by folder, but
// just to be safe.
if (position.empty()) {
rewriter.replaceOp(insertOp, source);
return success();
}

Value tileSlice = source;
Value sliceIndex = vector::getAsValues(rewriter, loc, position[0]).front();
if (position.size() == 2) {
// Two indices case: Insert single element into tile.
// We need to first extract the existing slice and update the element.
tileSlice = rewriter.create<arm_sme::MoveTileSliceToVectorOp>(
loc, insertOp.getDest(), sliceIndex);
tileSlice = rewriter.create<vector::InsertOp>(loc, source, tileSlice,
position[1]);
}

// Insert the slice into the destination tile.
rewriter.replaceOpWithNewOp<arm_sme::MoveVectorToTileSliceOp>(
insertOp, tileSlice, insertOp.getDest(), sliceIndex);
return success();
}
};

} // namespace

void mlir::populateVectorToArmSMEPatterns(RewritePatternSet &patterns,
Expand All @@ -581,5 +691,7 @@ void mlir::populateVectorToArmSMEPatterns(RewritePatternSet &patterns,
SplatOpToArmSMELowering, TransferReadToArmSMELowering,
TransferWriteToArmSMELowering, TransposeOpToArmSMELowering,
VectorLoadToArmSMELowering, VectorStoreToArmSMELowering,
VectorOuterProductToArmSMELowering>(&ctx);
VectorOuterProductToArmSMELowering,
VectorExtractToArmSMELowering, VectorInsertToArmSMELowering>(
&ctx);
}
115 changes: 1 addition & 114 deletions mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -520,118 +520,6 @@ struct OuterProductOpConversion
}
};

/// Lower `vector.extract` using `arm_sme.move_tile_slice_to_vector`.
///
/// Example:
/// ```
/// %el = vector.extract %tile[%row, %col]: i32 from vector<[4]x[4]xi32>
/// ```
/// Becomes:
/// ```
/// %slice = arm_sme.move_tile_slice_to_vector %tile[%row]
/// : vector<[4]xi32> from vector<[4]x[4]xi32>
/// %el = vector.extract %slice[%col] : i32 from vector<[4]xi32>
/// ```
struct VectorExtractToArmSMELowering
: public ConvertOpToLLVMPattern<vector::ExtractOp> {
using ConvertOpToLLVMPattern<vector::ExtractOp>::ConvertOpToLLVMPattern;

LogicalResult
matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
VectorType sourceType = extractOp.getSourceVectorType();
if (!isValidSMETileVectorType(sourceType))
return failure();

auto loc = extractOp.getLoc();
auto position = extractOp.getMixedPosition();

Value sourceVector = extractOp.getVector();

// Extract entire vector. Should be handled by folder, but just to be safe.
if (position.empty()) {
rewriter.replaceOp(extractOp, sourceVector);
return success();
}

Value sliceIndex = vector::getAsValues(rewriter, loc, position[0]).front();
auto moveTileSliceToVector =
rewriter.create<arm_sme::MoveTileSliceToVectorOp>(loc, sourceVector,
sliceIndex);

if (position.size() == 1) {
// Single index case: Extracts a 1D slice.
rewriter.replaceOp(extractOp, moveTileSliceToVector);
return success();
}

// Two indices case: Extracts a single element.
assert(position.size() == 2);
rewriter.replaceOpWithNewOp<vector::ExtractOp>(
extractOp, moveTileSliceToVector, position[1]);

return success();
}
};

/// Lower `vector.insert` using `arm_sme.move_vector_to_tile_slice` and
/// `arm_sme.move_tile_slice_to_vector`.
///
/// Example:
/// ```
/// %new_tile = vector.insert %el, %tile[%row, %col]
/// : i32 into vector<[4]x[4]xi32>
/// ```
/// Becomes:
/// ```
/// %slice = arm_sme.move_tile_slice_to_vector %tile[%row]
/// : vector<[4]xi32> from vector<[4]x[4]xi32>
/// %new_slice = vector.insert %el, %slice[%col] : i32 into vector<[4]xi32>
/// %new_tile = arm_sme.move_vector_to_tile_slice %new_slice, %tile, %row
/// : vector<[4]xi32> into vector<[4]x[4]xi32>
/// ```
struct VectorInsertToArmSMELowering
: public ConvertOpToLLVMPattern<vector::InsertOp> {
using ConvertOpToLLVMPattern<vector::InsertOp>::ConvertOpToLLVMPattern;

LogicalResult
matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
VectorType resultType = insertOp.getResult().getType();

if (!isValidSMETileVectorType(resultType))
return failure();

auto loc = insertOp.getLoc();
auto position = insertOp.getMixedPosition();

Value source = adaptor.getSource();

// Overwrite entire vector with value. Should be handled by folder, but
// just to be safe.
if (position.empty()) {
rewriter.replaceOp(insertOp, source);
return success();
}

Value tileSlice = source;
Value sliceIndex = vector::getAsValues(rewriter, loc, position[0]).front();
if (position.size() == 2) {
// Two indices case: Insert single element into tile.
// We need to first extract the existing slice and update the element.
tileSlice = rewriter.create<arm_sme::MoveTileSliceToVectorOp>(
loc, adaptor.getDest(), sliceIndex);
tileSlice = rewriter.create<vector::InsertOp>(loc, source, tileSlice,
position[1]);
}

// Insert the slice into the destination tile.
rewriter.replaceOpWithNewOp<arm_sme::MoveVectorToTileSliceOp>(
insertOp, tileSlice, adaptor.getDest(), sliceIndex);
return success();
}
};

} // namespace

void mlir::configureArmSMELegalizeForExportTarget(
Expand Down Expand Up @@ -661,6 +549,5 @@ void mlir::populateArmSMELegalizeForLLVMExportPatterns(
patterns.add<
LoadTileSliceToArmSMELowering, MoveTileSliceToVectorArmSMELowering,
MoveVectorToTileSliceToArmSMELowering, StoreTileSliceToArmSMELowering,
OuterProductOpConversion, ZeroOpConversion, VectorExtractToArmSMELowering,
VectorInsertToArmSMELowering>(converter);
OuterProductOpConversion, ZeroOpConversion>(converter);
}

0 comments on commit c4c52d4

Please sign in to comment.