[Mlir-commits] [mlir] a56a7d9 - [MLIR][OpenMP] Support schedule chunk size with various bit width
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jan 18 20:38:01 PST 2022
Author: Peixin-Qiao
Date: 2022-01-19T12:36:53+08:00
New Revision: a56a7d99e8594975cd783f10ed3c0ccbf9479229
URL: https://github.com/llvm/llvm-project/commit/a56a7d99e8594975cd783f10ed3c0ccbf9479229
DIFF: https://github.com/llvm/llvm-project/commit/a56a7d99e8594975cd783f10ed3c0ccbf9479229.diff
LOG: [MLIR][OpenMP] Support schedule chunk size with various bit width
The chunk size in schedule clause is one integer expression, which can
be either constant integer or integer variable. Fix schedule clause in
MLIR Op Def to support integer expression with different bit width.
Reviewed By: shraiysh
Differential Revision: https://reviews.llvm.org/D116073
Added:
Modified:
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
mlir/test/Dialect/OpenMP/ops.mlir
mlir/test/Target/LLVMIR/openmp-llvm.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 33cb9535edbb5..553775cd7faf2 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -295,7 +295,8 @@ verifyScheduleModifiers(OpAsmParser &parser,
static ParseResult
parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
SmallVectorImpl<SmallString<12>> &modifiers,
- Optional<OpAsmParser::OperandType> &chunkSize) {
+ Optional<OpAsmParser::OperandType> &chunkSize,
+ Type &chunkType) {
if (parser.parseLParen())
return failure();
@@ -307,7 +308,7 @@ parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
if (keyword == "static" || keyword == "dynamic" || keyword == "guided") {
if (succeeded(parser.parseOptionalEqual())) {
chunkSize = OpAsmParser::OperandType{};
- if (parser.parseOperand(*chunkSize))
+ if (parser.parseOperand(*chunkSize) || parser.parseColonType(chunkType))
return failure();
} else {
chunkSize = llvm::NoneType::None;
@@ -341,7 +342,7 @@ static void printScheduleClause(OpAsmPrinter &p, ClauseScheduleKind sched,
Value scheduleChunkVar) {
p << "schedule(" << stringifyClauseScheduleKind(sched).lower();
if (scheduleChunkVar)
- p << " = " << scheduleChunkVar;
+ p << " = " << scheduleChunkVar << " : " << scheduleChunkVar.getType();
if (modifier)
p << ", " << stringifyScheduleModifier(*modifier);
if (simd)
@@ -631,6 +632,7 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
SmallString<8> schedule;
SmallVector<SmallString<12>> modifiers;
Optional<OpAsmParser::OperandType> scheduleChunkSize;
+ Type scheduleChunkType;
// Compute the position of clauses in operand segments
int currPos = 0;
@@ -751,7 +753,8 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
clauseSegments[pos[linearClause] + 1] = linearSteps.size();
} else if (clauseKeyword == "schedule") {
if (checkAllowed(scheduleClause) ||
- parseScheduleClause(parser, schedule, modifiers, scheduleChunkSize))
+ parseScheduleClause(parser, schedule, modifiers, scheduleChunkSize,
+ scheduleChunkType))
return failure();
if (scheduleChunkSize) {
clauseSegments[pos[scheduleClause]] = 1;
@@ -906,10 +909,9 @@ static ParseResult parseClauses(OpAsmParser &parser, OperationState &result,
result.addAttribute("simd_modifier", attr);
}
}
- if (scheduleChunkSize) {
- auto chunkSizeType = parser.getBuilder().getI32Type();
- parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
- }
+ if (scheduleChunkSize)
+ parser.resolveOperand(*scheduleChunkSize, scheduleChunkType,
+ result.operands);
}
segments.insert(segments.end(), clauseSegments.begin(), clauseSegments.end());
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index c7a5b0846502e..14d38d44f6ab8 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -646,10 +646,20 @@ convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
// Find the loop configuration.
llvm::Value *step = moduleTranslation.lookupValue(loop.step()[0]);
llvm::Type *ivType = step->getType();
- llvm::Value *chunk =
- loop.schedule_chunk_var()
- ? moduleTranslation.lookupValue(loop.schedule_chunk_var())
- : llvm::ConstantInt::get(ivType, 1);
+ llvm::Value *chunk = nullptr;
+ if (loop.schedule_chunk_var()) {
+ llvm::Value *chunkVar =
+ moduleTranslation.lookupValue(loop.schedule_chunk_var());
+ llvm::Type *chunkVarType = chunkVar->getType();
+ assert(chunkVarType->isIntegerTy() &&
+ "chunk size must be one integer expression");
+ if (chunkVarType->getIntegerBitWidth() < ivType->getIntegerBitWidth())
+ chunk = builder.CreateSExt(chunkVar, ivType);
+ else if (chunkVarType->getIntegerBitWidth() > ivType->getIntegerBitWidth())
+ chunk = builder.CreateTrunc(chunkVar, ivType);
+ else
+ chunk = chunkVar;
+ }
SmallVector<omp::ReductionDeclareOp> reductionDecls;
collectReductionDecls(loop, reductionDecls);
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index fef4b8adef7bf..950f3d0d472a5 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -188,8 +188,7 @@ func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32
}
// CHECK-LABEL: omp_wsloop_pretty
-func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
- %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
+func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32, %chunk_var2 : i16) -> () {
// CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) private(%data_var : memref<i32>) collapse(2) ordered(2) {
@@ -201,24 +200,24 @@ func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
omp.yield
}
- // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}}) collapse(3) ordered(2)
+ // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}} : i32) collapse(3) ordered(2)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
- schedule(static = %chunk_var) collapse(3) {
+ schedule(static = %chunk_var : i32) collapse(3) {
omp.yield
}
- // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}, nonmonotonic) collapse(3) ordered(2)
+ // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i32, nonmonotonic) collapse(3) ordered(2)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
- schedule(dynamic = %chunk_var, nonmonotonic) collapse(3) {
+ schedule(dynamic = %chunk_var : i32, nonmonotonic) collapse(3) {
omp.yield
}
- // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}, monotonic) collapse(3) ordered(2)
+ // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i16, monotonic) collapse(3) ordered(2)
omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step) ordered(2) private(%data_var : memref<i32>)
firstprivate(%data_var : memref<i32>) lastprivate(%data_var : memref<i32>) linear(%data_var = %linear_var : memref<i32>)
- schedule(dynamic = %chunk_var, monotonic) collapse(3) {
+ schedule(dynamic = %chunk_var2 : i16, monotonic) collapse(3) {
omp.yield
}
diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
index 3b9d8f1c11651..94a0e0b3ae376 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
@@ -430,7 +430,24 @@ llvm.func @test_omp_wsloop_dynamic(%lb : i64, %ub : i64, %step : i64) -> () {
// CHECK: call void @__kmpc_dispatch_init_8u
// CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
// CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
- // CHECK br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+ // CHECK: br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+ llvm.call @body(%iv) : (i64) -> ()
+ omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
+llvm.func @body(i64)
+
+llvm.func @test_omp_wsloop_dynamic_chunk_const(%lb : i64, %ub : i64, %step : i64) -> () {
+ %chunk_size_const = llvm.mlir.constant(2 : i16) : i16
+ omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic = %chunk_size_const : i16) {
+ // CHECK: call void @__kmpc_dispatch_init_8u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 35, i64 {{.*}}, i64 %{{.*}}, i64 {{.*}}, i64 2)
+ // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u
+ // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+ // CHECK: br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
llvm.call @body(%iv) : (i64) -> ()
omp.yield
}
@@ -439,6 +456,62 @@ llvm.func @test_omp_wsloop_dynamic(%lb : i64, %ub : i64, %step : i64) -> () {
// -----
+llvm.func @body(i32)
+
+llvm.func @test_omp_wsloop_dynamic_chunk_var(%lb : i32, %ub : i32, %step : i32) -> () {
+ %1 = llvm.mlir.constant(1 : i64) : i64
+ %chunk_size_alloca = llvm.alloca %1 x i16 {bindc_name = "chunk_size", in_type = i16, uniq_name = "_QFsub1Echunk_size"} : (i64) -> !llvm.ptr<i16>
+ %chunk_size_var = llvm.load %chunk_size_alloca : !llvm.ptr<i16>
+ omp.wsloop (%iv) : i32 = (%lb) to (%ub) step (%step) schedule(dynamic = %chunk_size_var : i16) {
+ // CHECK: %[[CHUNK_SIZE:.*]] = sext i16 %{{.*}} to i32
+ // CHECK: call void @__kmpc_dispatch_init_4u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 35, i32 {{.*}}, i32 %{{.*}}, i32 {{.*}}, i32 %[[CHUNK_SIZE]])
+ // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_4u
+ // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+ // CHECK: br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+ llvm.call @body(%iv) : (i32) -> ()
+ omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
+llvm.func @body(i32)
+
+llvm.func @test_omp_wsloop_dynamic_chunk_var2(%lb : i32, %ub : i32, %step : i32) -> () {
+ %1 = llvm.mlir.constant(1 : i64) : i64
+ %chunk_size_alloca = llvm.alloca %1 x i64 {bindc_name = "chunk_size", in_type = i64, uniq_name = "_QFsub1Echunk_size"} : (i64) -> !llvm.ptr<i64>
+ %chunk_size_var = llvm.load %chunk_size_alloca : !llvm.ptr<i64>
+ omp.wsloop (%iv) : i32 = (%lb) to (%ub) step (%step) schedule(dynamic = %chunk_size_var : i64) {
+ // CHECK: %[[CHUNK_SIZE:.*]] = trunc i64 %{{.*}} to i32
+ // CHECK: call void @__kmpc_dispatch_init_4u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 35, i32 {{.*}}, i32 %{{.*}}, i32 {{.*}}, i32 %[[CHUNK_SIZE]])
+ // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_4u
+ // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+ // CHECK: br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+ llvm.call @body(%iv) : (i32) -> ()
+ omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
+llvm.func @body(i32)
+
+llvm.func @test_omp_wsloop_dynamic_chunk_var3(%lb : i32, %ub : i32, %step : i32, %chunk_size : i32) -> () {
+ omp.wsloop (%iv) : i32 = (%lb) to (%ub) step (%step) schedule(dynamic = %chunk_size : i32) {
+ // CHECK: call void @__kmpc_dispatch_init_4u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 35, i32 {{.*}}, i32 %{{.*}}, i32 {{.*}}, i32 %{{.*}})
+ // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_4u
+ // CHECK: %[[cond:.*]] = icmp ne i32 %[[continue]], 0
+ // CHECK: br i1 %[[cond]], label %omp_loop.header{{.*}}, label %omp_loop.exit{{.*}}
+ llvm.call @body(%iv) : (i32) -> ()
+ omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
llvm.func @body(i64)
llvm.func @test_omp_wsloop_auto(%lb : i64, %ub : i64, %step : i64) -> () {
More information about the Mlir-commits
mailing list