Skip to content

[mlir][Vector] Remove more special case uses for extractelement/insertelement #130166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 24, 2025
Merged

[mlir][Vector] Remove more special case uses for extractelement/insertelement #130166

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80c6dbe
[mlir][Vector] Remove more special case uses for extractelement/inser…
Groverkss Mar 6, 2025
c85a2bd
add comment
Groverkss Mar 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -718,6 +718,8 @@ def Vector_ExtractOp :
let results = (outs AnyType:$result);

let builders = [
// Builder to extract a scalar from a rank-0 vector.
OpBuilder<(ins "Value":$source)>,
OpBuilder<(ins "Value":$source, "int64_t":$position)>,
OpBuilder<(ins "Value":$source, "OpFoldResult":$position)>,
OpBuilder<(ins "Value":$source, "ArrayRef<int64_t>":$position)>,
Expand Down Expand Up @@ -913,6 +915,8 @@ def Vector_InsertOp :
let results = (outs AnyVectorOfAnyRank:$result);

let builders = [
// Builder to insert a scalar/rank-0 vector into a rank-0 vector.
OpBuilder<(ins "Value":$source, "Value":$dest)>,
OpBuilder<(ins "Value":$source, "Value":$dest, "int64_t":$position)>,
OpBuilder<(ins "Value":$source, "Value":$dest, "OpFoldResult":$position)>,
OpBuilder<(ins "Value":$source, "Value":$dest, "ArrayRef<int64_t>":$position)>,
Expand Down
32 changes: 18 additions & 14 deletions mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -560,11 +560,9 @@ struct ElideUnitDimsInMultiDimReduction
} else {
// This means we are reducing all the dimensions, and all reduction
// dimensions are of size 1. So a simple extraction would do.
SmallVector<int64_t> zeroIdx(shape.size(), 0);
if (mask)
mask = rewriter.create<vector::ExtractOp>(loc, mask, zeroIdx);
cast = rewriter.create<vector::ExtractOp>(loc, reductionOp.getSource(),
zeroIdx);
mask = rewriter.create<vector::ExtractOp>(loc, mask);
cast = rewriter.create<vector::ExtractOp>(loc, reductionOp.getSource());
}

Value result =
Expand Down Expand Up @@ -698,16 +696,9 @@ struct ElideSingleElementReduction : public OpRewritePattern<ReductionOp> {
return failure();

Location loc = reductionOp.getLoc();
Value result;
if (vectorType.getRank() == 0) {
if (mask)
mask = rewriter.create<ExtractElementOp>(loc, mask);
result = rewriter.create<ExtractElementOp>(loc, reductionOp.getVector());
} else {
if (mask)
mask = rewriter.create<ExtractOp>(loc, mask, 0);
result = rewriter.create<ExtractOp>(loc, reductionOp.getVector(), 0);
}
if (mask)
mask = rewriter.create<ExtractOp>(loc, mask);
Value result = rewriter.create<ExtractOp>(loc, reductionOp.getVector());

if (Value acc = reductionOp.getAcc())
result = vector::makeArithReduction(rewriter, loc, reductionOp.getKind(),
Expand Down Expand Up @@ -1294,6 +1285,12 @@ void ExtractOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
setResultRanges(getResult(), argRanges.front());
}

void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
Value source) {
auto vectorTy = cast<VectorType>(source.getType());
build(builder, result, source, SmallVector<int64_t>(vectorTy.getRank(), 0));
}

void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
Value source, int64_t position) {
build(builder, result, source, ArrayRef<int64_t>{position});
Expand Down Expand Up @@ -2916,6 +2913,13 @@ void vector::InsertOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
setResultRanges(getResult(), argRanges[0].rangeUnion(argRanges[1]));
}

void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
Value source, Value dest) {
auto vectorTy = cast<VectorType>(dest.getType());
build(builder, result, source, dest,
SmallVector<int64_t>(vectorTy.getRank(), 0));
}

void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
Value source, Value dest, int64_t position) {
build(builder, result, source, dest, ArrayRef<int64_t>{position});
Expand Down
6 changes: 1 addition & 5 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorBroadcast.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,11 +52,7 @@ class BroadcastOpLowering : public OpRewritePattern<vector::BroadcastOp> {

// Stretching scalar inside vector (e.g. vector<1xf32>) can use splat.
if (srcRank <= 1 && dstRank == 1) {
Value ext;
if (srcRank == 0)
ext = rewriter.create<vector::ExtractElementOp>(loc, op.getSource());
else
ext = rewriter.create<vector::ExtractOp>(loc, op.getSource(), 0);
Value ext = rewriter.create<vector::ExtractOp>(loc, op.getSource());
rewriter.replaceOpWithNewOp<vector::SplatOp>(op, dstType, ext);
return success();
}
Expand Down
22 changes: 3 additions & 19 deletions mlir/lib/Dialect/Vector/Transforms/LowerVectorShapeCast.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,25 +189,9 @@ class ShapeCastOpRewritePattern : public OpRewritePattern<vector::ShapeCastOp> {
incIdx(resIdx, resultVectorType);
}

Value extract;
if (srcRank == 0) {
// 0-D vector special case
assert(srcIdx.empty() && "Unexpected indices for 0-D vector");
extract = rewriter.create<vector::ExtractElementOp>(
loc, op.getSourceVectorType().getElementType(), op.getSource());
} else {
extract =
rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx);
}

if (resRank == 0) {
// 0-D vector special case
assert(resIdx.empty() && "Unexpected indices for 0-D vector");
result = rewriter.create<vector::InsertElementOp>(loc, extract, result);
} else {
result =
rewriter.create<vector::InsertOp>(loc, extract, result, resIdx);
}
Value extract =
rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx);
result = rewriter.create<vector::InsertOp>(loc, extract, result, resIdx);
}
rewriter.replaceOp(op, result);
return success();
Expand Down
13 changes: 2 additions & 11 deletions mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -929,17 +929,8 @@ class RewriteScalarWrite : public OpRewritePattern<vector::TransferWriteOp> {
if (!xferOp.getPermutationMap().isMinorIdentity())
return failure();
// Only float and integer element types are supported.
Value scalar;
if (vecType.getRank() == 0) {
// vector.extract does not support vector<f32> etc., so use
// vector.extractelement instead.
scalar = rewriter.create<vector::ExtractElementOp>(xferOp.getLoc(),
xferOp.getVector());
} else {
SmallVector<int64_t> pos(vecType.getRank(), 0);
scalar = rewriter.create<vector::ExtractOp>(xferOp.getLoc(),
xferOp.getVector(), pos);
}
Value scalar =
rewriter.create<vector::ExtractOp>(xferOp.getLoc(), xferOp.getVector());
// Construct a scalar store.
if (isa<MemRefType>(xferOp.getSource().getType())) {
rewriter.replaceOpWithNewOp<memref::StoreOp>(
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ func.func @broadcast_vec2d_from_vec0d(%arg0: vector<f32>) -> vector<3x2xf32> {
// CHECK: %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<f32> to vector<1xf32>
// CHECK: %[[T1:.*]] = ub.poison : vector<3x2xf32>
// CHECK: %[[T2:.*]] = builtin.unrealized_conversion_cast %[[T1]] : vector<3x2xf32> to !llvm.array<3 x vector<2xf32>>
// CHECK: %[[T4:.*]] = llvm.mlir.constant(0 : index) : i64
// CHECK: %[[T4:.*]] = llvm.mlir.constant(0 : i64) : i64
// CHECK: %[[T5:.*]] = llvm.extractelement %[[T0]][%[[T4]] : i64] : vector<1xf32>
// CHECK: %[[T6Insert:.*]] = llvm.insertelement %[[T5]]
// CHECK: %[[T6:.*]] = llvm.shufflevector %[[T6Insert]]
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Dialect/Vector/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2658,7 +2658,7 @@ func.func @fold_extractelement_of_broadcast(%f: f32) -> f32 {

// CHECK-LABEL: func.func @fold_0d_vector_reduction
func.func @fold_0d_vector_reduction(%arg0: vector<f32>) -> f32 {
// CHECK-NEXT: %[[RES:.*]] = vector.extractelement %arg{{.*}}[] : vector<f32>
// CHECK-NEXT: %[[RES:.*]] = vector.extract %arg{{.*}}[] : f32 from vector<f32>
// CHECK-NEXT: return %[[RES]] : f32
%0 = vector.reduction <add>, %arg0 : vector<f32> into f32
return %0 : f32
Expand Down
8 changes: 2 additions & 6 deletions mlir/test/Dialect/Vector/scalar-vector-transfer-to-memref.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,7 @@ func.func @tensor_transfer_read_0d(%t: tensor<?x?x?xf32>, %idx: index) -> f32 {

// CHECK-LABEL: func @transfer_write_0d(
// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>, %[[idx:.*]]: index, %[[f:.*]]: f32
// CHECK: %[[bc:.*]] = vector.broadcast %[[f]] : f32 to vector<f32>
// CHECK: %[[extract:.*]] = vector.extractelement %[[bc]][] : vector<f32>
// CHECK: memref.store %[[extract]], %[[m]][%[[idx]], %[[idx]], %[[idx]]]
// CHECK: memref.store %[[f]], %[[m]][%[[idx]], %[[idx]], %[[idx]]]
func.func @transfer_write_0d(%m: memref<?x?x?xf32>, %idx: index, %f: f32) {
%0 = vector.broadcast %f : f32 to vector<f32>
vector.transfer_write %0, %m[%idx, %idx, %idx] : vector<f32>, memref<?x?x?xf32>
Expand All @@ -69,9 +67,7 @@ func.func @transfer_write_1d(%m: memref<?x?x?xf32>, %idx: index, %f: f32) {

// CHECK-LABEL: func @tensor_transfer_write_0d(
// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[idx:.*]]: index, %[[f:.*]]: f32
// CHECK: %[[bc:.*]] = vector.broadcast %[[f]] : f32 to vector<f32>
// CHECK: %[[extract:.*]] = vector.extractelement %[[bc]][] : vector<f32>
// CHECK: %[[r:.*]] = tensor.insert %[[extract]] into %[[t]][%[[idx]], %[[idx]], %[[idx]]]
// CHECK: %[[r:.*]] = tensor.insert %[[f]] into %[[t]][%[[idx]], %[[idx]], %[[idx]]]
// CHECK: return %[[r]]
func.func @tensor_transfer_write_0d(%t: tensor<?x?x?xf32>, %idx: index, %f: f32) -> tensor<?x?x?xf32> {
%0 = vector.broadcast %f : f32 to vector<f32>
Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Dialect/Vector/vector-shape-cast-lowering-transforms.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ func.func @shape_cast_1d3d(%arg0 : vector<6xf32>) -> vector<2x1x3xf32> {
// CHECK-LABEL: func.func @shape_cast_0d1d(
// CHECK-SAME: %[[ARG0:.*]]: vector<f32>) -> vector<1xf32> {
// CHECK: %[[UB:.*]] = ub.poison : vector<1xf32>
// CHECK: %[[EXTRACT0:.*]] = vector.extractelement %[[ARG0]][] : vector<f32>
// CHECK: %[[EXTRACT0:.*]] = vector.extract %[[ARG0]][] : f32 from vector<f32>
// CHECK: %[[RES:.*]] = vector.insert %[[EXTRACT0]], %[[UB]] [0] : f32 into vector<1xf32>
// CHECK: return %[[RES]] : vector<1xf32>
// CHECK: }
Expand All @@ -131,7 +131,7 @@ func.func @shape_cast_0d1d(%arg0 : vector<f32>) -> vector<1xf32> {
// CHECK-SAME: %[[ARG0:.*]]: vector<1xf32>) -> vector<f32> {
// CHECK: %[[UB:.*]] = ub.poison : vector<f32>
// CHECK: %[[EXTRACT0:.*]] = vector.extract %[[ARG0]][0] : f32 from vector<1xf32>
// CHECK: %[[RES:.*]] = vector.insertelement %[[EXTRACT0]], %[[UB]][] : vector<f32>
// CHECK: %[[RES:.*]] = vector.insert %[[EXTRACT0]], %[[UB]] [] : f32 into vector<f32>
// CHECK: return %[[RES]] : vector<f32>
// CHECK: }

Expand Down
Loading