Skip to content

[mlir][vector][affine] Allow --affine-super-vectorize to vectorize maxnumf/minnumf (2) #138730

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][vector][affine] Allow --affine-super-vectorize to vectorize maxnumf/minnumf (2) #138730

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 2 additions & 0 deletions mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ static Value getSupportedReduction(AffineForOp forOp, unsigned pos,
[](arith::MinimumFOp) { return arith::AtomicRMWKind::minimumf; })
.Case(
[](arith::MaximumFOp) { return arith::AtomicRMWKind::maximumf; })
.Case([](arith::MinNumFOp) { return arith::AtomicRMWKind::minnumf; })
.Case([](arith::MaxNumFOp) { return arith::AtomicRMWKind::maxnumf; })
.Case([](arith::MinSIOp) { return arith::AtomicRMWKind::mins; })
.Case([](arith::MaxSIOp) { return arith::AtomicRMWKind::maxs; })
.Case([](arith::MinUIOp) { return arith::AtomicRMWKind::minu; })
Expand Down
6 changes: 3 additions & 3 deletions mlir/lib/Dialect/Affine/IR/AffineOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3949,8 +3949,9 @@ static bool isResultTypeMatchAtomicRMWKind(Type resultType,
case arith::AtomicRMWKind::muli:
return isa<IntegerType>(resultType);
case arith::AtomicRMWKind::maximumf:
return isa<FloatType>(resultType);
case arith::AtomicRMWKind::minimumf:
case arith::AtomicRMWKind::maxnumf:
case arith::AtomicRMWKind::minnumf:
return isa<FloatType>(resultType);
case arith::AtomicRMWKind::maxs: {
auto intType = llvm::dyn_cast<IntegerType>(resultType);
Expand All @@ -3972,9 +3973,8 @@ static bool isResultTypeMatchAtomicRMWKind(Type resultType,
return isa<IntegerType>(resultType);
case arith::AtomicRMWKind::andi:
return isa<IntegerType>(resultType);
default:
return false;
}
llvm_unreachable("exhaustive switch");
}

LogicalResult AffineParallelOp::verify() {
Expand Down
16 changes: 11 additions & 5 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 @@ -690,12 +690,18 @@ Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
case arith::AtomicRMWKind::ori:
return builder.create<vector::ReductionOp>(vector.getLoc(),
CombiningKind::OR, vector);
// TODO: Add remaining reduction operations.
default:
(void)emitOptionalError(loc, "Reduction operation type not supported");
break;
case arith::AtomicRMWKind::maxnumf:
return builder.create<vector::ReductionOp>(vector.getLoc(),
CombiningKind::MAXNUMF, vector);
case arith::AtomicRMWKind::minnumf:
return builder.create<vector::ReductionOp>(vector.getLoc(),
CombiningKind::MINNUMF, vector);
case arith::AtomicRMWKind::assign:
(void)emitOptionalError(loc,
"Reduction operation type not supported (assign)");
return nullptr;
}
return nullptr;
llvm_unreachable("exhaustive switch");
}

std::optional<SmallVector<int64_t, 4>> ReductionOp::getShapeForUnroll() {
Expand Down
54 changes: 54 additions & 0 deletions mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,60 @@ func.func @vecdim_reduction_maxf(%in: memref<256x512xf32>, %out: memref<256xf32>

// -----

func.func @vecdim_reduction_minnumf(%in: memref<256x512xf32>, %out: memref<256xf32>) {
%cst = arith.constant 0x7FC00000 : f32
affine.for %i = 0 to 256 {
%final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
%ld = affine.load %in[%i, %j] : memref<256x512xf32>
%min = arith.minnumf %red_iter, %ld : f32
affine.yield %min : f32
}
affine.store %final_red, %out[%i] : memref<256xf32>
}
return
}

// CHECK-LABEL: @vecdim_reduction_minnumf
// CHECK: affine.for %{{.*}} = 0 to 256 {
// CHECK: %[[vmax:.*]] = arith.constant dense<0x7FC00000> : vector<128xf32>
// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xf32>) {
// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK: %[[min:.*]] = arith.minnumf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK: affine.yield %[[min]] : vector<128xf32>
// CHECK: }
// CHECK: %[[final_min:.*]] = vector.reduction <minnumf>, %[[vred:.*]] : vector<128xf32> into f32
// CHECK: affine.store %[[final_min]], %{{.*}} : memref<256xf32>
// CHECK: }

// -----

func.func @vecdim_reduction_maxnumf(%in: memref<256x512xf32>, %out: memref<256xf32>) {
%cst = arith.constant 0xFFC00000 : f32
affine.for %i = 0 to 256 {
%final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
%ld = affine.load %in[%i, %j] : memref<256x512xf32>
%max = arith.maxnumf %red_iter, %ld : f32
affine.yield %max : f32
}
affine.store %final_red, %out[%i] : memref<256xf32>
}
return
}

// CHECK-LABEL: @vecdim_reduction_maxnumf
// CHECK: affine.for %{{.*}} = 0 to 256 {
// CHECK: %[[vmin:.*]] = arith.constant dense<0xFFC00000> : vector<128xf32>
// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xf32>) {
// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK: %[[max:.*]] = arith.maxnumf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK: affine.yield %[[max]] : vector<128xf32>
// CHECK: }
// CHECK: %[[final_max:.*]] = vector.reduction <maxnumf>, %[[vred:.*]] : vector<128xf32> into f32
// CHECK: affine.store %[[final_max]], %{{.*}} : memref<256xf32>
// CHECK: }

// -----

func.func @vecdim_reduction_minsi(%in: memref<256x512xi32>, %out: memref<256xi32>) {
%cst = arith.constant 2147483647 : i32
affine.for %i = 0 to 256 {
Expand Down