[OMPIRBuilder] Fix static schedule in worksharing-loop directive

[peixin]

Summary

If no chunk_size is specified, OMPScheduleType for kmpc_for_static_init
call is Static and the chunk size is set to 1. If chunk_size is specified,
OMPScheduleType for kmpc_for_static_init call is StaticChunked. Also
fix lowering to LLVM IR for schdule clause in worksharing-loop directive.

[Leporacanthicus]

Generally looks good to me. Just one nit-pick and one "I don't really understand how this actually produces the correct result" (may be that I'm just missing something!)

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

726

I may be a bit confused here, but isn't this setting chunk to not-null, and then expecting the apply function to see it as null?

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

457

I'd like to see this having a separate name - it makes it so much easier to search the source code when the tests have distinct names! :)

[peixin]

Change the variable name"const_i32" to "static_chunk_size" and remove redundant code.

[peixin]

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

726

This code is the "else" branch, which is for non-static schedule. As OpenMP 5.1 [2.11.4] states, the chunk_size defaults to 1 when it is unspecified for dynamic or guided schedule. The chunk size cannot be specified for runtime or auto schedule. But in implementation, llvm openmp runtime uses kmpc_dispatch_init with the argument "chunk_size" of 1 for runtime or auto schedule.

To be short, schedule(static) is NOT the same as schedule(static,1).
schedule(dynamic) is the same as schedule(dynamic,1).
schedule(guided) is the same as schedule(guided,1).
schedule(runtime, 1) and schedule(auto, 1) is not allowed, but schedule(runtime) and schedule(auto) behave like schedule(runtime, 1) and schedule(auto, 1) in IR level.

I noticed "if (!chunk) chunk = one;" is also added in applyDynamicWorkshareLoop in OMPIRBuilder, so remove this code here.

[Meinersbur]

@peixin is correct. However, a "dispatch" loop is needed in addition to a "chunk" loop, please see TODO in the description of applyStaticWorkshareLoop. This patch would only iterate over chunk_size*omp_get_num_threads() iterations and leave any additional loop iterations unprocessed.

Implementation of static schedule without chunk-size: D114413

[peixin]

@Meinersbur Thanks for the clarification. Can you also handle the fix in mlir part in D114413 (initilize chunk to nullptr and add some test cases in oepnmp-llvm.mlir)? I will close this PR since it is not complete fix.

[Leporacanthicus]

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

726

Thanks, I obviously got confused by the "hidden" bit of code, and thought it was just continuation of the static wsloop code.

[Meinersbur]

In D114331#3148040, @peixin wrote:

@Meinersbur Thanks for the clarification. Can you also handle the fix in mlir part in D114413 (initilize chunk to nullptr and add some test cases in oepnmp-llvm.mlir)? I will close this PR since it is not complete fix.

See update to D114413

[peixin]

Thanks @Meinersbur .

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OpenMPIRBuilder::applyStaticWorkshareLoop(DebugLoc DL, CanonicalLoopInfo *CLI,
   Builder.SetInsertPoint(CLI->getPreheader()->getTerminator());
   Constant *Zero = ConstantInt::get(IVTy, 0);
   Constant *One = ConstantInt::get(IVTy, 1);
   Builder.CreateStore(Zero, PLowerBound);
   Value *UpperBound = Builder.CreateSub(CLI->getTripCount(), One);
   Builder.CreateStore(UpperBound, PUpperBound);
   Builder.CreateStore(One, PStride);
 
-  // FIXME: schedule(static) is NOT the same as schedule(static,1)
-  if (!Chunk)
+  Constant *SchedulingType;
+  if (Chunk) {
+    SchedulingType = ConstantInt::get(
+        I32Type, static_cast<int>(OMPScheduleType::StaticChunked));
+  } else {
+    SchedulingType =
+        ConstantInt::get(I32Type, static_cast<int>(OMPScheduleType::Static));
     Chunk = One;
+  }
 
   Value *ThreadNum = getOrCreateThreadID(SrcLoc);
 
-  Constant *SchedulingType =
-      ConstantInt::get(I32Type, static_cast<int>(OMPScheduleType::Static));
-
   // Call the "init" function and update the trip count of the loop with the
   // value it produced.
   Builder.CreateCall(StaticInit,
                      {SrcLoc, ThreadNum, SchedulingType, PLastIter, PLowerBound,
                       PUpperBound, PStride, One, Chunk});
   Value *LowerBound = Builder.CreateLoad(IVTy, PLowerBound);
   Value *InclusiveUpperBound = Builder.CreateLoad(IVTy, PUpperBound);
   Value *TripCountMinusOne = Builder.CreateSub(InclusiveUpperBound, LowerBound);

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

convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
   omp::ClauseScheduleKind schedule = omp::ClauseScheduleKind::Static;
   if (loop.schedule_val().hasValue())
     schedule =
         *omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue());
 
   // 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())
+    chunk = moduleTranslation.lookupValue(loop.schedule_chunk_var());
 
   SmallVector<omp::ReductionDeclareOp> reductionDecls;
   collectReductionDecls(loop, reductionDecls);
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
       findAllocaInsertPoint(builder, moduleTranslation);
 
   // Allocate space for privatized reduction variables.
   SmallVector<llvm::Value *> privateReductionVariables;
if (loop.schedule_modifier().hasValue()) {
         break;
       default:
         // Nothing to do here.
         break;
       }
     }
     afterIP = ompBuilder->applyDynamicWorkshareLoop(
         ompLoc.DL, loopInfo, allocaIP, schedType, !loop.nowait(), chunk);
   }
 
   // Continue building IR after the loop. Note that the LoopInfo returned by
   // `collapseLoops` points inside the outermost loop and is intended for
   // potential further loop transformations. Use the insertion point stored
   // before collapsing loops instead.
   builder.restoreIP(afterIP);
 
   // Process the reductions if required.

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

^bb0(%arg1: i64):
     llvm.store %3, %4 : !llvm.ptr<f32>
     omp.yield
   }) {inclusive, operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0, 0]> : vector<10xi32>} : (i64, i64, i64) -> ()
   llvm.return
 }
 
 // -----
 
+llvm.func @body(i32)
+
+llvm.func @test_omp_wsloop_static_1(%lb : i32, %ub : i32, %step : i32) -> () {
+ omp.wsloop (%iv) : i32 = (%lb) to (%ub) step (%step) schedule(static) {
+   // CHECK: call void @__kmpc_for_static_init_4u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 34, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32 1, i32 1)
+   // CHECK: call void @__kmpc_for_static_fini
+   llvm.call @body(%iv) : (i32) -> ()
+   omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
+llvm.func @body(i32)
+
+llvm.func @test_omp_wsloop_static_2(%lb : i32, %ub : i32, %step : i32) -> () {
+ %static_chunk_size = llvm.mlir.constant(1 : i32) : i32
+ omp.wsloop (%iv) : i32 = (%lb) to (%ub) step (%step) schedule(static = %static_chunk_size) {
+   // CHECK: call void @__kmpc_for_static_init_4u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 33, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32 1, i32 1)
+   // CHECK: call void @__kmpc_for_static_fini
+   llvm.call @body(%iv) : (i32) -> ()
+   omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
+llvm.func @body(i32)
+
+llvm.func @test_omp_wsloop_static_2(%lb : i32, %ub : i32, %step : i32) -> () {
+ %static_chunk_size = llvm.mlir.constant(2 : i32) : i32
+ omp.wsloop (%iv) : i32 = (%lb) to (%ub) step (%step) schedule(static = %static_chunk_size) {
+   // CHECK: call void @__kmpc_for_static_init_4u(%struct.ident_t* @{{.*}}, i32 %{{.*}}, i32 33, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32* %{{.*}}, i32 1, i32 2)
+   // CHECK: call void @__kmpc_for_static_fini
+   llvm.call @body(%iv) : (i32) -> ()
+   omp.yield
+ }
+ llvm.return
+}
+
+// -----
+
 llvm.func @body(i64)
 
 llvm.func @test_omp_wsloop_dynamic(%lb : i64, %ub : i64, %step : i64) -> () {
  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic) {
   // 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{{.*}}