Skip to content

Commit d3eb65f

Browse files
authored
[MLIR][OpenMP] Lowering aligned clause to LLVM IR for SIMD directive (llvm#119536)
This patch, - Added a translation support for aligned clause in SIMD directive by passing the alignment details to "llvm.assume" intrinsic. - Updated the insertion point for llvm.assume intrinsic call in "OMPIRBuilder.cpp". - Added a check in aligned clause MLIR lowering, to ensure that the alignment value must be a power of 2.
1 parent e3ec5a7 commit d3eb65f

File tree

7 files changed

+89
-25
lines changed

7 files changed

+89
-25
lines changed

clang/test/OpenMP/irbuilder_simd_aligned.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,11 @@ void simple(float *a, float *b, int *c) {
7070
// CHECK-NEXT: br label [[FOR_COND]], !llvm.loop [[LOOP3:![0-9]+]]
7171
// CHECK: for.end:
7272
// CHECK-NEXT: [[TMP4:%.*]] = load ptr, ptr [[A_ADDR]], align 8
73+
// CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[TMP4]], i64 128) ]
7374
// CHECK-NEXT: [[TMP5:%.*]] = load ptr, ptr [[P]], align 8
75+
// CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[TMP5]], i64 64) ]
7476
// CHECK-NEXT: [[ARRAYDECAY:%.*]] = getelementptr inbounds [32 x i32], ptr [[D]], i64 0, i64 0
77+
// CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[ARRAYDECAY]], i64 16) ]
7578
// CHECK-NEXT: store i32 3, ptr [[I1]], align 4
7679
// CHECK-NEXT: [[TMP6:%.*]] = getelementptr inbounds nuw [[STRUCT_ANON]], ptr [[AGG_CAPTURED]], i32 0, i32 0
7780
// CHECK-NEXT: store ptr [[I1]], ptr [[TMP6]], align 8
@@ -82,9 +85,6 @@ void simple(float *a, float *b, int *c) {
8285
// CHECK-NEXT: [[DOTCOUNT:%.*]] = load i32, ptr [[DOTCOUNT_ADDR]], align 4
8386
// CHECK-NEXT: br label [[OMP_LOOP_PREHEADER:%.*]]
8487
// CHECK: omp_loop.preheader:
85-
// CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[TMP4]], i64 128) ]
86-
// CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[TMP5]], i64 64) ]
87-
// CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[ARRAYDECAY]], i64 16) ]
8888
// CHECK-NEXT: br label [[OMP_LOOP_HEADER:%.*]]
8989
// CHECK: omp_loop.header:
9090
// CHECK-NEXT: [[OMP_LOOP_IV:%.*]] = phi i32 [ 0, [[OMP_LOOP_PREHEADER]] ], [ [[OMP_LOOP_NEXT:%.*]], [[OMP_LOOP_INC:%.*]] ]

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -613,6 +613,8 @@ addAlignedClause(lower::AbstractConverter &converter,
613613
// Do not generate alignment assumption if alignment is less than or equal to
614614
// 0.
615615
if (alignment > 0) {
616+
// alignment value must be power of 2
617+
assert((alignment & (alignment - 1)) == 0 && "alignment is not power of 2");
616618
auto &objects = std::get<omp::ObjectList>(clause.t);
617619
if (!objects.empty())
618620
genObjectList(objects, converter, alignedVars);

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5302,10 +5302,11 @@ void OpenMPIRBuilder::applySimd(CanonicalLoopInfo *CanonicalLoop,
53025302
Loop *L = LI.getLoopFor(CanonicalLoop->getHeader());
53035303
if (AlignedVars.size()) {
53045304
InsertPointTy IP = Builder.saveIP();
5305-
Builder.SetInsertPoint(CanonicalLoop->getPreheader()->getTerminator());
53065305
for (auto &AlignedItem : AlignedVars) {
53075306
Value *AlignedPtr = AlignedItem.first;
53085307
Value *Alignment = AlignedItem.second;
5308+
Instruction *loadInst = dyn_cast<Instruction>(AlignedPtr);
5309+
Builder.SetInsertPoint(loadInst->getNextNode());
53095310
Builder.CreateAlignmentAssumption(F->getDataLayout(),
53105311
AlignedPtr, Alignment);
53115312
}

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1993,6 +1993,7 @@ TEST_F(OpenMPIRBuilderTest, ApplySimdCustomAligned) {
19931993
OpenMPIRBuilder OMPBuilder(*M);
19941994
IRBuilder<> Builder(BB);
19951995
const int AlignmentValue = 32;
1996+
llvm::BasicBlock *sourceBlock = Builder.GetInsertBlock();
19961997
AllocaInst *Alloc1 =
19971998
Builder.CreateAlloca(Builder.getPtrTy(), Builder.getInt64(1));
19981999
LoadInst *Load1 = Builder.CreateLoad(Alloc1->getAllocatedType(), Alloc1);
@@ -2031,13 +2032,12 @@ TEST_F(OpenMPIRBuilderTest, ApplySimdCustomAligned) {
20312032

20322033
// Check if number of assumption instructions is equal to number of aligned
20332034
// variables
2034-
BasicBlock *LoopPreheader = CLI->getPreheader();
2035-
size_t NumAssummptionCallsInPreheader = count_if(
2036-
*LoopPreheader, [](Instruction &I) { return isa<AssumeInst>(I); });
2035+
size_t NumAssummptionCallsInPreheader =
2036+
count_if(*sourceBlock, [](Instruction &I) { return isa<AssumeInst>(I); });
20372037
EXPECT_EQ(NumAssummptionCallsInPreheader, AlignedVars.size());
20382038

20392039
// Check if variables are correctly aligned
2040-
for (Instruction &Instr : *LoopPreheader) {
2040+
for (Instruction &Instr : *sourceBlock) {
20412041
if (!isa<AssumeInst>(Instr))
20422042
continue;
20432043
AssumeInst *AssumeInstruction = cast<AssumeInst>(&Instr);

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -150,10 +150,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
150150
<< " operation";
151151
};
152152

153-
auto checkAligned = [&todo](auto op, LogicalResult &result) {
154-
if (!op.getAlignedVars().empty() || op.getAlignments())
155-
result = todo("aligned");
156-
};
157153
auto checkAllocate = [&todo](auto op, LogicalResult &result) {
158154
if (!op.getAllocateVars().empty() || !op.getAllocatorVars().empty())
159155
result = todo("allocate");
@@ -275,7 +271,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
275271
})
276272
.Case([&](omp::ParallelOp op) { checkAllocate(op, result); })
277273
.Case([&](omp::SimdOp op) {
278-
checkAligned(op, result);
279274
checkLinear(op, result);
280275
checkNontemporal(op, result);
281276
checkPrivate(op, result);
@@ -2302,6 +2297,24 @@ convertOmpSimd(Operation &opInst, llvm::IRBuilderBase &builder,
23022297

23032298
llvm::MapVector<llvm::Value *, llvm::Value *> alignedVars;
23042299
llvm::omp::OrderKind order = convertOrderKind(simdOp.getOrder());
2300+
llvm::BasicBlock *sourceBlock = builder.GetInsertBlock();
2301+
std::optional<ArrayAttr> alignmentValues = simdOp.getAlignments();
2302+
mlir::OperandRange operands = simdOp.getAlignedVars();
2303+
for (size_t i = 0; i < operands.size(); ++i) {
2304+
llvm::Value *alignment = nullptr;
2305+
llvm::Value *llvmVal = moduleTranslation.lookupValue(operands[i]);
2306+
llvm::Type *ty = llvmVal->getType();
2307+
if (auto intAttr = llvm::dyn_cast<IntegerAttr>((*alignmentValues)[i])) {
2308+
alignment = builder.getInt64(intAttr.getInt());
2309+
assert(ty->isPointerTy() && "Invalid type for aligned variable");
2310+
assert(alignment && "Invalid alignment value");
2311+
auto curInsert = builder.saveIP();
2312+
builder.SetInsertPoint(sourceBlock->getTerminator());
2313+
llvmVal = builder.CreateLoad(ty, llvmVal);
2314+
builder.restoreIP(curInsert);
2315+
alignedVars[llvmVal] = alignment;
2316+
}
2317+
}
23052318
ompBuilder->applySimd(loopInfo, alignedVars,
23062319
simdOp.getIfExpr()
23072320
? moduleTranslation.lookupValue(simdOp.getIfExpr())
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
2+
3+
//CHECK-LABEL: define void @_QPsimd_aligned_pointer() {
4+
//CHECK: %[[A_PTR:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8 }, i64 1, align 8
5+
//CHECK: %[[A_VAL:.*]] = load ptr, ptr %[[A_PTR]], align 8
6+
//CHECK: call void @llvm.assume(i1 true) [ "align"(ptr %[[A_VAL]], i64 256) ]
7+
llvm.func @_QPsimd_aligned_pointer() {
8+
%1 = llvm.mlir.constant(1 : i64) : i64
9+
%2 = llvm.alloca %1 x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)> {bindc_name = "x"} : (i64) -> !llvm.ptr
10+
%3 = llvm.alloca %1 x i32 {bindc_name = "i", pinned} : (i64) -> !llvm.ptr
11+
%4 = llvm.mlir.constant(1 : i32) : i32
12+
%5 = llvm.mlir.constant(10 : i32) : i32
13+
%6 = llvm.mlir.constant(1 : i32) : i32
14+
omp.simd aligned(%2 : !llvm.ptr -> 256 : i64) {
15+
omp.loop_nest (%arg0) : i32 = (%4) to (%5) inclusive step (%6) {
16+
llvm.store %arg0, %3 : i32, !llvm.ptr
17+
omp.yield
18+
}
19+
}
20+
llvm.return
21+
}
22+
23+
//CHECK-LABEL: define void @_QPsimd_aligned_cptr() {
24+
//CHECK: %[[A_CPTR:.*]] = alloca %_QM__fortran_builtinsT__builtin_c_ptr, i64 1, align 8
25+
//CHECK: %[[A_VAL:.*]] = load ptr, ptr %[[A_CPTR]], align 8
26+
//CHECK: call void @llvm.assume(i1 true) [ "align"(ptr %[[A_VAL]], i64 256) ]
27+
llvm.func @_QPsimd_aligned_cptr() {
28+
%0 = llvm.mlir.constant(1 : i64) : i64
29+
%1 = llvm.alloca %0 x !llvm.struct<"_QM__fortran_builtinsT__builtin_c_ptr", (i64)> {bindc_name = "a"} : (i64) -> !llvm.ptr
30+
%2 = llvm.mlir.constant(1 : i64) : i64
31+
%3 = llvm.alloca %2 x i32 {bindc_name = "i", pinned} : (i64) -> !llvm.ptr
32+
%4 = llvm.mlir.constant(1 : i32) : i32
33+
%5 = llvm.mlir.constant(10 : i32) : i32
34+
%6 = llvm.mlir.constant(1 : i32) : i32
35+
omp.simd aligned(%1 : !llvm.ptr -> 256 : i64) {
36+
omp.loop_nest (%arg0) : i32 = (%4) to (%5) inclusive step (%6) {
37+
llvm.store %arg0, %3 : i32, !llvm.ptr
38+
omp.yield
39+
}
40+
}
41+
llvm.return
42+
}
43+
44+
//CHECK-LABEL: define void @_QPsimd_aligned_allocatable() {
45+
//CHECK: %[[A_ADDR:.*]] = alloca { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, i64 1, align 8
46+
//CHECK: %[[A_VAL:.*]] = load ptr, ptr %[[A_ADDR]], align 8
47+
//CHECK: call void @llvm.assume(i1 true) [ "align"(ptr %[[A_VAL]], i64 256) ]
48+
llvm.func @_QPsimd_aligned_allocatable() {
49+
%0 = llvm.mlir.constant(1 : i64) : i64
50+
%1 = llvm.alloca %0 x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> {bindc_name = "a"} : (i64) -> !llvm.ptr
51+
%2 = llvm.mlir.constant(1 : i32) : i32
52+
%3 = llvm.mlir.constant(10 : i32) : i32
53+
%4 = llvm.mlir.constant(1 : i32) : i32
54+
omp.simd aligned(%1 : !llvm.ptr -> 256 : i64) {
55+
omp.loop_nest (%arg0) : i32 = (%2) to (%3) inclusive step (%4) {
56+
omp.yield
57+
}
58+
}
59+
llvm.return
60+
}

mlir/test/Target/LLVMIR/openmp-todo.mlir

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -127,18 +127,6 @@ llvm.func @sections_private(%x : !llvm.ptr) {
127127
llvm.return
128128
}
129129

130-
// -----
131-
132-
llvm.func @simd_aligned(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr) {
133-
// expected-error@below {{not yet implemented: Unhandled clause aligned in omp.simd operation}}
134-
// expected-error@below {{LLVM Translation failed for operation: omp.simd}}
135-
omp.simd aligned(%x : !llvm.ptr -> 32) {
136-
omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
137-
omp.yield
138-
}
139-
}
140-
llvm.return
141-
}
142130

143131
// -----
144132

0 commit comments

Comments
 (0)