Skip to content

Commit 3728f65

Browse files
joviliastjoviliast
committed
Allow 16 bit floating point operand for LLVM_AtomicRMWOp
As far as AMDGPU target supports vectorization for atomic_rmw operation, allow construction of LLVM_AtomicRMWOp with 16 bit floating point values. This patch enables building of LLVM_AtomicRMWOp with fixed vectors of 16 bit fp values as operands. See also: #94845, #95393, #95394 Signed-off-by: Ilya Veselov <[email protected]>
1 parent 2d3119c commit 3728f65

File tree

6 files changed

+57
-25
lines changed

6 files changed

+57
-25
lines changed

‎mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,16 @@ class LLVM_VectorOf<Type element> : Type<
139139
class LLVM_ScalarOrVectorOf<Type element> :
140140
AnyTypeOf<[element, LLVM_VectorOf<element>]>;
141141

142+
// Type constraint accepting an LLVM fixed vector type with an additional constraint
143+
// on the vector element type.
144+
class LLVM_FixedVectorOf<Type element> : Type<
145+
And<[LLVM_AnyFixedVector.predicate,
146+
SubstLeaves<
147+
"$_self",
148+
"::mlir::LLVM::getVectorElementType($_self)",
149+
element.predicate>]>,
150+
"LLVM dialect-compatible fixed vector of " # element.summary>;
151+
142152
// Base class for LLVM operations. Defines the interface to the llvm::IRBuilder
143153
// used to translate to proper LLVM IR and the interface to the mlir::OpBuilder
144154
// used to import from LLVM IR.

‎mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1737,7 +1737,8 @@ def LLVM_ConstantOp
17371737
// Atomic operations.
17381738
//
17391739

1740-
def LLVM_AtomicRMWType : AnyTypeOf<[LLVM_AnyFloat, LLVM_AnyPointer, AnySignlessInteger]>;
1740+
def LLVM_AtomicRMWType
1741+
: AnyTypeOf<[LLVM_AnyFloat, LLVM_AnyPointer, AnySignlessInteger, LLVM_FixedVectorOf<LLVM_AnyFloat>]>;
17411742

17421743
def LLVM_AtomicRMWOp : LLVM_MemAccessOpBase<"atomicrmw", [
17431744
TypesMatchWith<"result #0 and operand #1 have the same type",

‎mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3008,9 +3008,17 @@ void AtomicRMWOp::build(OpBuilder &builder, OperationState &state,
30083008

30093009
LogicalResult AtomicRMWOp::verify() {
30103010
auto valType = getVal().getType();
3011-
if (getBinOp() == AtomicBinOp::fadd || getBinOp() == AtomicBinOp::fsub ||
3012-
getBinOp() == AtomicBinOp::fmin || getBinOp() == AtomicBinOp::fmax) {
3013-
if (!mlir::LLVM::isCompatibleFloatingPointType(valType))
3011+
if (getBinOp() == AtomicBinOp::fadd && isCompatibleVectorType(valType)) {
3012+
Type elemType = getVectorElementType(valType);
3013+
// Only 16 bit floating point elements are supported for now.
3014+
if (!(isCompatibleFloatingPointType(elemType) &&
3015+
elemType.getIntOrFloatBitWidth() == 16))
3016+
return emitOpError("unexpected LLVM IR type for vector element");
3017+
} else if (getBinOp() == AtomicBinOp::fadd ||
3018+
getBinOp() == AtomicBinOp::fsub ||
3019+
getBinOp() == AtomicBinOp::fmin ||
3020+
getBinOp() == AtomicBinOp::fmax) {
3021+
if (!isCompatibleFloatingPointType(valType))
30143022
return emitOpError("expected LLVM IR floating point type");
30153023
} else if (getBinOp() == AtomicBinOp::xchg) {
30163024
DataLayout dataLayout = DataLayout::closest(*this);

‎mlir/test/Dialect/LLVMIR/invalid.mlir

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -643,6 +643,14 @@ func.func @atomicrmw_expected_float(%i32_ptr : !llvm.ptr, %i32 : i32) {
643643

644644
// -----
645645

646+
func.func @atomicrmw_unexpected_vector_element(%ptr : !llvm.ptr, %f32_vec : vector<3xf32>) {
647+
// expected-error@+1 {{unexpected LLVM IR type for vector element}}
648+
%0 = llvm.atomicrmw fadd %ptr, %f32_vec unordered : !llvm.ptr, vector<3xf32>
649+
llvm.return
650+
}
651+
652+
// -----
653+
646654
func.func @atomicrmw_unexpected_xchg_type(%i1_ptr : !llvm.ptr, %i1 : i1) {
647655
// expected-error@+1 {{unexpected LLVM IR type for 'xchg' bin_op}}
648656
%0 = llvm.atomicrmw xchg %i1_ptr, %i1 unordered : !llvm.ptr, i1

‎mlir/test/Dialect/LLVMIR/roundtrip.mlir

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -420,11 +420,13 @@ func.func @atomic_store(%val : f32, %large_val : i256, %ptr : !llvm.ptr) {
420420
}
421421

422422
// CHECK-LABEL: @atomicrmw
423-
func.func @atomicrmw(%ptr : !llvm.ptr, %val : f32) {
423+
func.func @atomicrmw(%ptr : !llvm.ptr, %val : f32, %f16_vec : vector<2xf16>) {
424424
// CHECK: llvm.atomicrmw fadd %{{.*}}, %{{.*}} monotonic : !llvm.ptr, f32
425425
%0 = llvm.atomicrmw fadd %ptr, %val monotonic : !llvm.ptr, f32
426+
// CHECK: llvm.atomicrmw fadd %{{.*}}, %{{.*}} monotonic : !llvm.ptr, vector<2xf16>
427+
%1 = llvm.atomicrmw fadd %ptr, %f16_vec monotonic : !llvm.ptr, vector<2xf16>
426428
// CHECK: llvm.atomicrmw volatile fsub %{{.*}}, %{{.*}} syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr, f32
427-
%1 = llvm.atomicrmw volatile fsub %ptr, %val syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr, f32
429+
%2 = llvm.atomicrmw volatile fsub %ptr, %val syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr, f32
428430
llvm.return
429431
}
430432

‎mlir/test/Target/LLVMIR/llvmir.mlir

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1496,50 +1496,53 @@ llvm.func @elements_constant_3d_array() -> !llvm.array<2 x array<2 x array<2 x i
14961496
// CHECK-LABEL: @atomicrmw
14971497
llvm.func @atomicrmw(
14981498
%f32_ptr : !llvm.ptr, %f32 : f32,
1499+
%f16_vec_ptr : !llvm.ptr, %f16_vec : vector<2xf16>,
14991500
%i32_ptr : !llvm.ptr, %i32 : i32) {
15001501
// CHECK: atomicrmw fadd ptr %{{.*}}, float %{{.*}} monotonic
15011502
%0 = llvm.atomicrmw fadd %f32_ptr, %f32 monotonic : !llvm.ptr, f32
1503+
// CHECK: atomicrmw fadd ptr %{{.*}}, <2 x half> %{{.*}} monotonic
1504+
%1 = llvm.atomicrmw fadd %f16_vec_ptr, %f16_vec monotonic : !llvm.ptr, vector<2xf16>
15021505
// CHECK: atomicrmw fsub ptr %{{.*}}, float %{{.*}} monotonic
1503-
%1 = llvm.atomicrmw fsub %f32_ptr, %f32 monotonic : !llvm.ptr, f32
1506+
%2 = llvm.atomicrmw fsub %f32_ptr, %f32 monotonic : !llvm.ptr, f32
15041507
// CHECK: atomicrmw fmax ptr %{{.*}}, float %{{.*}} monotonic
1505-
%2 = llvm.atomicrmw fmax %f32_ptr, %f32 monotonic : !llvm.ptr, f32
1508+
%3 = llvm.atomicrmw fmax %f32_ptr, %f32 monotonic : !llvm.ptr, f32
15061509
// CHECK: atomicrmw fmin ptr %{{.*}}, float %{{.*}} monotonic
1507-
%3 = llvm.atomicrmw fmin %f32_ptr, %f32 monotonic : !llvm.ptr, f32
1510+
%4 = llvm.atomicrmw fmin %f32_ptr, %f32 monotonic : !llvm.ptr, f32
15081511
// CHECK: atomicrmw xchg ptr %{{.*}}, float %{{.*}} monotonic
1509-
%4 = llvm.atomicrmw xchg %f32_ptr, %f32 monotonic : !llvm.ptr, f32
1512+
%5 = llvm.atomicrmw xchg %f32_ptr, %f32 monotonic : !llvm.ptr, f32
15101513
// CHECK: atomicrmw add ptr %{{.*}}, i32 %{{.*}} acquire
1511-
%5 = llvm.atomicrmw add %i32_ptr, %i32 acquire : !llvm.ptr, i32
1514+
%6 = llvm.atomicrmw add %i32_ptr, %i32 acquire : !llvm.ptr, i32
15121515
// CHECK: atomicrmw sub ptr %{{.*}}, i32 %{{.*}} release
1513-
%6 = llvm.atomicrmw sub %i32_ptr, %i32 release : !llvm.ptr, i32
1516+
%7 = llvm.atomicrmw sub %i32_ptr, %i32 release : !llvm.ptr, i32
15141517
// CHECK: atomicrmw and ptr %{{.*}}, i32 %{{.*}} acq_rel
1515-
%7 = llvm.atomicrmw _and %i32_ptr, %i32 acq_rel : !llvm.ptr, i32
1518+
%8 = llvm.atomicrmw _and %i32_ptr, %i32 acq_rel : !llvm.ptr, i32
15161519
// CHECK: atomicrmw nand ptr %{{.*}}, i32 %{{.*}} seq_cst
1517-
%8 = llvm.atomicrmw nand %i32_ptr, %i32 seq_cst : !llvm.ptr, i32
1520+
%9 = llvm.atomicrmw nand %i32_ptr, %i32 seq_cst : !llvm.ptr, i32
15181521
// CHECK: atomicrmw or ptr %{{.*}}, i32 %{{.*}} monotonic
1519-
%9 = llvm.atomicrmw _or %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1522+
%10 = llvm.atomicrmw _or %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15201523
// CHECK: atomicrmw xor ptr %{{.*}}, i32 %{{.*}} monotonic
1521-
%10 = llvm.atomicrmw _xor %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1524+
%11 = llvm.atomicrmw _xor %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15221525
// CHECK: atomicrmw max ptr %{{.*}}, i32 %{{.*}} monotonic
1523-
%11 = llvm.atomicrmw max %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1526+
%12 = llvm.atomicrmw max %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15241527
// CHECK: atomicrmw min ptr %{{.*}}, i32 %{{.*}} monotonic
1525-
%12 = llvm.atomicrmw min %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1528+
%13 = llvm.atomicrmw min %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15261529
// CHECK: atomicrmw umax ptr %{{.*}}, i32 %{{.*}} monotonic
1527-
%13 = llvm.atomicrmw umax %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1530+
%14 = llvm.atomicrmw umax %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15281531
// CHECK: atomicrmw umin ptr %{{.*}}, i32 %{{.*}} monotonic
1529-
%14 = llvm.atomicrmw umin %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1532+
%15 = llvm.atomicrmw umin %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15301533
// CHECK: atomicrmw uinc_wrap ptr %{{.*}}, i32 %{{.*}} monotonic
1531-
%15 = llvm.atomicrmw uinc_wrap %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1534+
%16 = llvm.atomicrmw uinc_wrap %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15321535
// CHECK: atomicrmw udec_wrap ptr %{{.*}}, i32 %{{.*}} monotonic
1533-
%16 = llvm.atomicrmw udec_wrap %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1536+
%17 = llvm.atomicrmw udec_wrap %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15341537
// CHECK: atomicrmw usub_cond ptr %{{.*}}, i32 %{{.*}} monotonic
1535-
%17 = llvm.atomicrmw usub_cond %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1538+
%18 = llvm.atomicrmw usub_cond %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15361539
// CHECK: atomicrmw usub_sat ptr %{{.*}}, i32 %{{.*}} monotonic
1537-
%18 = llvm.atomicrmw usub_sat %i32_ptr, %i32 monotonic : !llvm.ptr, i32
1540+
%19 = llvm.atomicrmw usub_sat %i32_ptr, %i32 monotonic : !llvm.ptr, i32
15381541

15391542
// CHECK: atomicrmw volatile
15401543
// CHECK-SAME: syncscope("singlethread")
15411544
// CHECK-SAME: align 8
1542-
%19 = llvm.atomicrmw volatile udec_wrap %i32_ptr, %i32 syncscope("singlethread") monotonic {alignment = 8 : i64} : !llvm.ptr, i32
1545+
%20 = llvm.atomicrmw volatile udec_wrap %i32_ptr, %i32 syncscope("singlethread") monotonic {alignment = 8 : i64} : !llvm.ptr, i32
15431546
llvm.return
15441547
}
15451548

0 commit comments

Comments
 (0)