[mlir][OpenMP] Add nontemporal clause definition to simd construct

[domada]

Summary

simd nontemporal construct is represented as a list of variables which have low locality accross simd iterations.

Added verifier of nontemporal clause. MLIR tests were updated to test correctness of MLIR definition of nontemporal clause.

[kiranchandramohan]

Thanks for this patch. Have a question regarding the custom printer/parser.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

398

439

Why do we need a custom printer or parser? I see that the Flush op takes in a Variadic<OpenMP_PointerLikeType> and does not need a custom one.

[domada]

Applied remarks - removed custom parser and printer of nontemporal clause.

[domada]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

439

Hi Kiran,
thanks for your review. I added custom parser and printer of nontemporal clause in the first patch, because I wanted to obtain similar format to other OpenMP clauses.

The format of the nontemporal clause was as follows:
nontemporal (var1 : type_of_var1, var2 :type_of_var2)

Without custom parser and printer the format is as follows:
nontemporal (var1, var2 : type_of_var1, type_of_var2)

Please let me know if the new format is ok for you.

[kiranchandramohan]

LGTM. The new format looks OK. The earlier one is probably more readable. At some point we will have to look into whether we can make all of this uniform. Hopefully, at that point there will be support for interleaving operands and types in a list (assuming this is not already there, have you checked?).

flang/lib/Lower/OpenMP.cpp

 
 static void genOMP(Fortran::lower::AbstractConverter &converter,
                    Fortran::lower::pft::Evaluation &eval,
                    const Fortran::parser::OpenMPLoopConstruct &loopConstruct) {
 
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
   mlir::Location currentLocation = converter.getCurrentLocation();
   llvm::SmallVector<mlir::Value> lowerBound, upperBound, step, linearVars,
-      linearStepVars, reductionVars, alignedVars;
+      linearStepVars, reductionVars, alignedVars, nontemporalVars;
   mlir::Value scheduleChunkClauseOperand, ifClauseOperand;
   mlir::Attribute scheduleClauseOperand, noWaitClauseOperand,
       orderedClauseOperand, orderClauseOperand;
   mlir::IntegerAttr simdlenClauseOperand, safelenClauseOperand;
   SmallVector<Attribute> reductionDeclSymbols;
   Fortran::lower::StatementContext stmtCtx;
   const auto &loopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
       std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
static void genOMP(Fortran::lower::AbstractConverter &converter,
   }
 
   // 2.9.3.1 SIMD construct
   // TODO: Support all the clauses
   if (llvm::omp::OMPD_simd == ompDirective) {
     TypeRange resultType;
     auto SimdLoopOp = firOpBuilder.create<mlir::omp::SimdLoopOp>(
         currentLocation, resultType, lowerBound, upperBound, step, alignedVars,
-        nullptr, ifClauseOperand,
+        nullptr, ifClauseOperand, nontemporalVars,
         orderClauseOperand.dyn_cast_or_null<omp::ClauseOrderKindAttr>(),
         simdlenClauseOperand, safelenClauseOperand,
         /*inclusive=*/firOpBuilder.getUnitAttr());
     createBodyOfOp<omp::SimdLoopOp>(SimdLoopOp, converter, currentLocation,
                                     eval, &loopOpClauseList, iv);
     return;
   }
 

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

let description = [{
     The `alignment_values` attribute additionally specifies alignment of each
     corresponding aligned operand. Note that `$aligned_vars` and
     `alignment_values` should contain the same number of elements.
 
     When an if clause is present and evaluates to false, the preferred number of
     iterations to be executed concurrently is one, regardless of whether
     a simdlen clause is specified.
 
+    The optional `nontemporal` attribute specifies variables which have low
+    temporal locality across the iterations where they are accessed.
+
     The optional `order` attribute specifies which order the iterations of the
     associate loops are executed in. Currently the only option for this
     attribute is "concurrent".
 
     When a simdlen clause is present, the preferred number of iterations to be
     executed concurrently is the value provided to the simdlen clause.
 
     The safelen clause specifies that no two concurrent iterations within a
let summary = "simd loop construct";
 
   // TODO: Add other clauses
   let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              Variadic<IntLikeType>:$upperBound,
              Variadic<IntLikeType>:$step,
              Variadic<OpenMP_PointerLikeType>:$aligned_vars,
              OptionalAttr<I64ArrayAttr>:$alignment_values,
              Optional<I1>:$if_expr,
+             Variadic<OpenMP_PointerLikeType>:$nontemporal_vars,
              OptionalAttr<OrderKindAttr>:$order_val,
              ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$simdlen,
              ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$safelen,
              UnitAttr:$inclusive
      );
 
   let regions = (region AnyRegion:$region);
   let assemblyFormat = [{
     oilist(`aligned` `(`
               custom<AlignedClause>($aligned_vars, type($aligned_vars),
                                    $alignment_values) `)`
           |`if` `(` $if_expr `)`
+          |`nontemporal` `(`  $nontemporal_vars `:` type($nontemporal_vars) `)`
           |`order` `(` custom<ClauseAttr>($order_val) `)`
           |`simdlen` `(` $simdlen  `)`
           |`safelen` `(` $safelen  `)`
     ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
                                   type($step), $inclusive) attr-dict
   }];
 
   let extraClassDeclaration = [{

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

for (unsigned i = 0; i < linearVarsSize; ++i) {
     p << linearVars[i];
     if (linearStepVars.size() > i)
       p << " = " << linearStepVars[i];
     p << " : " << linearVars[i].getType() << separator;
   }
 }
 
 //===----------------------------------------------------------------------===//
+// Verifier for Nontemporal Clause
+//===----------------------------------------------------------------------===//
+
+static LogicalResult
+verifyNontemporalClause(Operation *op, OperandRange nontemporalVariables) {
+
+  // Check if each var is unique - OpenMP 5.0 -> 2.9.3.1 section
+  DenseSet<Value> nontemporalItems;
+  for (const auto &it : nontemporalVariables)
+    if (!nontemporalItems.insert(it).second)
+      return op->emitOpError() << "nontemporal variable used more than once";
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // Parser, verifier and printer for Aligned Clause
 //===----------------------------------------------------------------------===//
 static LogicalResult
 verifyAlignedClause(Operation *op, std::optional<ArrayAttr> alignmentValues,
                     OperandRange alignedVariables) {
   // Check if number of alignment values equals to number of aligned variables
   if (!alignedVariables.empty()) {
     if (!alignmentValues || alignmentValues->size() != alignedVariables.size())
if (this->getLowerBound().empty()) {
     return emitOpError() << "empty lowerbound for simd loop operation";
   }
   if (this->getSimdlen().has_value() && this->getSafelen().has_value() &&
       this->getSimdlen().value() > this->getSafelen().value()) {
     return emitOpError()
            << "simdlen clause and safelen clause are both present, but the "
               "simdlen value is not less than or equal to safelen value";
   }
-  return verifyAlignedClause(*this, this->getAlignmentValues(),
-                             this->getAlignedVars());
+  if (verifyAlignedClause(*this, this->getAlignmentValues(),
+                          this->getAlignedVars())
+          .failed())
+    return failure();
+  if (verifyNontemporalClause(*this, this->getNontemporalVars()).failed())
+    return failure();
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReductionOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseAtomicReductionRegion(OpAsmParser &parser,
                                               Region &region) {

mlir/test/Dialect/OpenMP/invalid.mlir

 
 // -----
 
 func.func @omp_simdloop(%lb : index, %ub : index, %step : i32) -> () {
   // expected-error @below {{op failed to verify that all of {lowerBound, upperBound, step} have same type}}
   "omp.simdloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = array<i32: 1,1,1,0,0>} :
+  }) {operand_segment_sizes = array<i32: 1,1,1,0,0,0>} :
     (index, index, i32) -> ()
 
   return
 }
 
 // -----
 
 func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index,

 func.func @omp_simdloop_aligned_mismatch(%arg0 : index, %arg1 : index,
                                          %arg2 : index, %arg3 : memref<i32>,
                                          %arg4 : memref<i32>) -> () {
   //  expected-error @below {{op expected as many alignment values as aligned variables}}
   "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg4) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
   }) {alignment_values = [128],
-      operand_segment_sizes = array<i32: 1, 1, 1, 2, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
+      operand_segment_sizes = array<i32: 1, 1, 1, 2, 0, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
   return
 }
 
 // -----
 
 func.func @omp_simdloop_aligned_negative(%arg0 : index, %arg1 : index,
                                          %arg2 : index, %arg3 : memref<i32>,
                                          %arg4 : memref<i32>) -> () {
   //  expected-error @below {{op alignment should be greater than 0}}
   "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg4) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
-  }) {alignment_values = [-1, 128], operand_segment_sizes = array<i32: 1, 1, 1,2, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
+  }) {alignment_values = [-1, 128], operand_segment_sizes = array<i32: 1, 1, 1,2, 0, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
   return
 }
 
 // -----
 
 func.func @omp_simdloop_unexpected_alignment(%arg0 : index, %arg1 : index,
                                              %arg2 : index, %arg3 : memref<i32>,
                                              %arg4 : memref<i32>) -> () {
   //  expected-error @below {{unexpected alignment values attribute}}
   "omp.simdloop"(%arg0, %arg1, %arg2) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
-  }) {alignment_values = [1, 128], operand_segment_sizes = array<i32: 1, 1, 1, 0, 0>} : (index, index, index) -> ()
+  }) {alignment_values = [1, 128], operand_segment_sizes = array<i32: 1, 1, 1, 0, 0, 0>} : (index, index, index) -> ()
   return
 }
 
 // -----
 
 func.func @omp_simdloop_aligned_float(%arg0 : index, %arg1 : index,
                                       %arg2 : index, %arg3 : memref<i32>,
                                       %arg4 : memref<i32>) -> () {
   //  expected-error @below {{failed to satisfy constraint: 64-bit integer array attribute}}
   "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg4) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
-  }) {alignment_values = [1.5, 128], operand_segment_sizes = array<i32: 1, 1, 1,2, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
+  }) {alignment_values = [1.5, 128], operand_segment_sizes = array<i32: 1, 1, 1,2, 0, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
   return
 }
 
 // -----
 
 func.func @omp_simdloop_aligned_the_same_var(%arg0 : index, %arg1 : index,
                                              %arg2 : index, %arg3 : memref<i32>,
                                              %arg4 : memref<i32>) -> () {
   //  expected-error @below {{aligned variable used more than once}}
   "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg3) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
-  }) {alignment_values = [1, 128], operand_segment_sizes = array<i32: 1, 1, 1,2, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
+  }) {alignment_values = [1, 128], operand_segment_sizes = array<i32: 1, 1, 1,2, 0, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
+  return
+}
+
+// -----
+
+func.func @omp_simdloop_nontemporal_the_same_var(%arg0 : index,
+                                                 %arg1 : index,
+                                                 %arg2 : index,
+                                                 %arg3 : memref<i32>) -> () {
+  //  expected-error @below {{nontemporal variable used more than once}}
+  "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg3) ({
+    ^bb0(%arg5: index):
+      "omp.yield"() : () -> ()
+  }) {operand_segment_sizes = array<i32: 1, 1, 1, 0, 0, 2>} : (index, index, index, memref<i32>, memref<i32>) -> ()
   return
 }
 
 // -----
 
 func.func @omp_simdloop_order_value(%lb : index, %ub : index, %step : index) {
   // expected-error @below {{invalid clause value: 'default'}}
   omp.simdloop order(default) for (%iv): index = (%lb) to (%ub) step (%step) {

mlir/test/Dialect/OpenMP/ops.mlir

 }
 
 // CHECK-LABEL: omp_simdloop
 func.func @omp_simdloop(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.simdloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = array<i32: 1,1,1,0,0>} :
+  }) {operand_segment_sizes = array<i32: 1,1,1,0,0,0>} :
     (index, index, index) -> ()
 
   return
 }
 
 // CHECK-LABEL: omp_simdloop_aligned_list
 func.func @omp_simdloop_aligned_list(%arg0 : index, %arg1 : index, %arg2 : index,
                                      %arg3 : memref<i32>, %arg4 : memref<i32>) -> () {
   // CHECK:      omp.simdloop   aligned(%{{.*}} : memref<i32> -> 32 : i64,
   // CHECK-SAME: %{{.*}} : memref<i32> -> 128 : i64)
   // CHECK-SAME: for  (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
   "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg4) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
   }) {alignment_values = [32, 128],
-      operand_segment_sizes = array<i32: 1, 1, 1,2, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
+      operand_segment_sizes = array<i32: 1, 1, 1, 2, 0, 0>} : (index, index, index, memref<i32>, memref<i32>) -> ()
   return
 }
 
 // CHECK-LABEL: omp_simdloop_aligned_single
 func.func @omp_simdloop_aligned_single(%arg0 : index, %arg1 : index, %arg2 : index,
                                        %arg3 : memref<i32>, %arg4 : memref<i32>) -> () {
   // CHECK:      omp.simdloop   aligned(%{{.*}} : memref<i32> -> 32 : i64)
   // CHECK-SAME: for  (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
   "omp.simdloop"(%arg0, %arg1, %arg2, %arg3) ({
     ^bb0(%arg5: index):
       "omp.yield"() : () -> ()
   }) {alignment_values = [32],
-      operand_segment_sizes = array<i32: 1, 1, 1,1, 0>} : (index, index, index, memref<i32>) -> ()
+      operand_segment_sizes = array<i32: 1, 1, 1, 1, 0, 0>} : (index, index, index, memref<i32>) -> ()
+  return
+}
+
+// CHECK-LABEL: omp_simdloop_nontemporal_list
+func.func @omp_simdloop_nontemporal_list(%arg0 : index,
+                                         %arg1 : index,
+                                         %arg2 : index,
+                                         %arg3 : memref<i32>,
+                                         %arg4 : memref<i64>) -> () {
+  // CHECK:      omp.simdloop   nontemporal(%{{.*}}, %{{.*}} : memref<i32>, memref<i64>)
+  // CHECK-SAME: for  (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
+  "omp.simdloop"(%arg0, %arg1, %arg2, %arg3, %arg4) ({
+    ^bb0(%arg5: index):
+      "omp.yield"() : () -> ()
+  }) {operand_segment_sizes = array<i32: 1, 1, 1, 0, 0, 2>} : (index, index, index, memref<i32>, memref<i64>) -> ()
+  return
+}
+
+// CHECK-LABEL: omp_simdloop_nontemporal_single
+func.func @omp_simdloop_nontemporal_single(%arg0 : index,
+                                           %arg1 : index,
+                                           %arg2 : index,
+                                           %arg3 : memref<i32>,
+                                           %arg4 : memref<i64>) -> () {
+  // CHECK:      omp.simdloop   nontemporal(%{{.*}} : memref<i32>)
+  // CHECK-SAME: for  (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
+  "omp.simdloop"(%arg0, %arg1, %arg2, %arg3) ({
+    ^bb0(%arg5: index):
+      "omp.yield"() : () -> ()
+  }) {operand_segment_sizes = array<i32: 1, 1, 1, 0, 0, 1>} : (index, index, index, memref<i32>) -> ()
   return
 }
 
 // CHECK-LABEL: omp_simdloop_pretty
 func.func @omp_simdloop_pretty(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   omp.simdloop for (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield

 func.func @omp_simdloop_pretty_if(%lb : index, %ub : index, %step : index, %if_cond : i1) -> () {
   // CHECK: omp.simdloop if(%{{.*}}) for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   omp.simdloop if(%if_cond) for (%iv): index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
   return
 }
 
+// CHECK-LABEL:   func.func @omp_simdloop_pretty_nontemporal
+func.func @omp_simdloop_pretty_nontemporal(%lb : index,
+                                           %ub : index,
+                                           %step : index,
+                                           %data_var : memref<i32>,
+                                           %data_var1 : memref<i32>) -> () {
+  // CHECK:      omp.simdloop   nontemporal(%{{.*}}, %{{.*}} : memref<i32>, memref<i32>)
+  // CHECK-SAME: for  (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
+  omp.simdloop nontemporal(%data_var, %data_var1 : memref<i32>, memref<i32>)
+    for (%iv) : index = (%lb) to (%ub) step (%step) {
+      omp.yield
+  }
+  return
+}
 // CHECK-LABEL: omp_simdloop_pretty_order
 func.func @omp_simdloop_pretty_order(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop order(concurrent)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   omp.simdloop order(concurrent) for (%iv): index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
   return

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

"omp.simdloop" (%lb, %ub, %step) ({
     ^bb0(%iv: i64):
       %3 = llvm.mlir.constant(2.000000e+00 : f32) : f32
       // The form of the emitted IR is controlled by OpenMPIRBuilder and
       // tested there. Just check that the right metadata is added.
       // CHECK: llvm.access.group
       %4 = llvm.getelementptr %arg0[%iv] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
       llvm.store %3, %4 : !llvm.ptr<f32>
       omp.yield
-  }) {operand_segment_sizes = array<i32: 1,1,1,0,0>} :
+  }) {operand_segment_sizes = array<i32: 1,1,1,0,0,0>} :
     (i64, i64, i64) -> ()
 
   llvm.return
 }
 // CHECK: llvm.loop.parallel_accesses
 // CHECK-NEXT: llvm.loop.vectorize.enable
 
 // -----