[OpenMP][flang]Lower NUM_THREADS clause for parallel construct

Author: SouraVX

flang/lib/Lower/OpenMP.cpp

@@ -15,6 +15,7 @@
 #include "flang/Lower/FIRBuilder.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Parser/parse-tree.h"
+#include "flang/Semantics/tools.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
@@ -87,19 +88,34 @@ genOMP(Fortran::lower::AbstractConverter &absConv,
     auto &firOpBuilder = absConv.getFirOpBuilder();
     auto currentLocation = absConv.getCurrentLocation();
     auto insertPt = firOpBuilder.saveInsertionPoint();
+
+    // Clauses.
+    // FIXME: Add support for other clauses.
+    mlir::Value numThreads;
+
+    const auto &parallelOpClauseList =
+        std::get<Fortran::parser::OmpClauseList>(blockDirective.t);
+    for (const auto &clause : parallelOpClauseList.v) {
+      if (const auto &numThreadsClause =
+              std::get_if<Fortran::parser::OmpClause::NumThreads>(&clause.u)) {
+        // OMPIRBuilder expects `NUM_THREAD` clause as a `Value`.
+        numThreads = absConv.genExprValue(
+            *Fortran::semantics::GetExpr(numThreadsClause->v));
+      }
+    }
     llvm::ArrayRef<mlir::Type> argTy;
-    mlir::ValueRange range;
-    llvm::SmallVector<int32_t, 6> operandSegmentSizes(6 /*Size=*/,
-                                                      0 /*Value=*/);
-    // create and insert the operation.
+    Attribute defaultValue, procBindValue;
+    // Create and insert the operation.
+    // Create the Op with empty ranges for clauses that are yet to be lowered.
     auto parallelOp = firOpBuilder.create<mlir::omp::ParallelOp>(
-        currentLocation, argTy, range);
-    parallelOp.setAttr(mlir::omp::ParallelOp::getOperandSegmentSizeAttr(),
-                       firOpBuilder.getI32VectorAttr(operandSegmentSizes));
-    parallelOp.getRegion().push_back(new Block{});
+        currentLocation, argTy, Value(), numThreads,
+        defaultValue.dyn_cast_or_null<StringAttr>(), ValueRange(), ValueRange(),
+        ValueRange(), ValueRange(),
+        procBindValue.dyn_cast_or_null<StringAttr>());
+    firOpBuilder.createBlock(&parallelOp.getRegion());
     auto &block = parallelOp.getRegion().back();
     firOpBuilder.setInsertionPointToStart(&block);
-    // ensure the block is well-formed.
+    // Ensure the block is well-formed.
     firOpBuilder.create<mlir::omp::TerminatorOp>(currentLocation);
     firOpBuilder.restoreInsertionPoint(insertPt);
   }

flang/test/Lower/OpenMP/empty-omp-parallel.f90

@@ -10,13 +10,15 @@
 program parallel
 
         integer :: a,b,c
+        integer :: num_threads
 ! This and last statements are just for the sake ensuring that the
 ! operation is created/inserted correctly and does not break/interfere with
 ! other pieces which may be present before/after the operation.
 ! However this test does not verify operation corresponding to this
 ! statment.
         c = a + b
 !$OMP PARALLEL
+!$OMP END PARALLEL
 !FIRDialect: omp.parallel {
 !FIRDialect-NEXT: omp.terminator
 !FIRDialect-NEXT: }
@@ -25,9 +27,35 @@ program parallel
 !LLVMIRDialect-NEXT: omp.terminator
 !LLVMIRDialect-NEXT: }
 
+!$OMP PARALLEL NUM_THREADS(16)
+!$OMP END PARALLEL
+        num_threads = 4
+!$OMP PARALLEL NUM_THREADS(num_threads)
+!$OMP END PARALLEL
+
+!FIRDialect: omp.parallel num_threads(%{{.*}} : i32) {
+!FIRDialect-NEXT: omp.terminator
+!FIRDialect-NEXT: }
+
+!LLVMIRDialect: omp.parallel num_threads(%{{.*}} : !llvm.i32) {
+!LLVMIRDialect-NEXT: omp.terminator
+!LLVMIRDialect-NEXT: }
+
+
+!LLVMIR-LABEL: call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{.*}})
 !LLVMIR: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN:.*]] to {{.*}}
+
+!LLVMIR: %[[GLOBAL_THREAD_NUM1:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{.*}})
+!LLVMIR: call void @__kmpc_push_num_threads(%struct.ident_t* @{{.*}}, i32 %[[GLOBAL_THREAD_NUM1]], i32 16)
+!LLVMIR: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN1:.*]] to {{.*}}
+
+!LLVMIR: %[[GLOBAL_THREAD_NUM2:.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @{{.*}})
+!LLVMIR: call void @__kmpc_push_num_threads(%struct.ident_t* @{{.*}}, i32 %[[GLOBAL_THREAD_NUM2]], i32 %{{.*}})
+!LLVMIR: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN2:.*]] to {{.*}}
+
+!LLVMIR: define internal void @[[OMP_OUTLINED_FN2]]
+!LLVMIR: define internal void @[[OMP_OUTLINED_FN1]]
 !LLVMIR: define internal void @[[OMP_OUTLINED_FN]]
-!$OMP END PARALLEL
         b = a + c
 
 end program