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

[domada]

Summary

simd aligned construct is represented as pair of variable which needs to be aligned and corresponding alignment value.

Added parser, printer and verifier of aligned clause. MLIR tests were updated to test correctness of MLIR definition of aligned clause.

[kiranchandramohan]

Thanks for this patch @domada. Have a few questions/comments.

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

389–391

Can you add this to a verifier?

415

If this is a variable then it is best to have the type OpenMP_PointerLikeType.

416

Does the standard specify I32?
Also, if this is a constant then it should be an attribute, something like

ConfinedAttr<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$aligned_val,

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

145

OpAsmParser has a parseCommaSeparatedList, can that be used instead?

mlir/test/Dialect/OpenMP/ops.mlir

383

Add a test with a list.

[domada]

Addressed review comments:

1. Modified the way how aligned clause is represented. Use I64ArrayAttr instead of Variadic to represent alignment values
2. Added verifier which checks if aligned clause was properly mapped
3. Added more tests MLIR tests

[domada]

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

416

OpenMP standard does not specify the size of alignment.

OpenMP standard allows to use multiple aligned clause to specify the alignment of different variables. That's why I needed to use an I64ArrayAttr to specify the alignment of each variable. I added to the verifier checks for positivity of every array element.

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

145

The format of alignment has changed. It is not variadic list any more. That's why I leave linear clauses unmodified.

I also use parseCommaSeparatedList to parse aligned clause.

[kiranchandramohan]

Looks mostly fine. Please add tests for all verification errors.

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

179

Nit:

186

Nit: Can you add a test for this?

194

Nit: Can you add a test for this?
For now, this is probably fine. But it is probably possible that there are multiple pointer vals pointing to the same address and this check will not catch it.

205

Nit: Can you add a test for this?

[domada]

Applied remarks. Added more tests.

[domada]

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

186

Added test case: omp_simdloop_unexpected_alignment

194

Added test case:
omp_simdloop_aligned_the_same_var

I checked that clang does not inform the user that two items point to the same address. The code given below: compiles without any warning: (build command: trunk_build/clang -fopenmp -c code_below.c )
int foo(int *A, int *B) {
int *p = B;
#pragma omp simd aligned(B) aligned(A, p :256)
for (int i = 0; i< 10; ++i) {
A[i] = i+44;
B[i] = i+2 + *p;
}
return 0;
}

I agree that we need to warn user about such buggy code, but in my opinion we need to add pointer aliasing analysis before OpenMP MLIR verification. I think that this issue should be addressed by different patch. Currently I can add TODO comment to remember about pointer aliaising.

205

Added test case:
omp_simdloop_aligned_float .
Alignment values are modelled as I64ArrayAttr in MLIR dialect. That's why I was not able to trigger "expected integer alignment" error. I am able to trigger generic MLIR parser error which corresponds to parsing error of I64ArrayAttr .

[kiranchandramohan]

LGTM.

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

194

OK, that is fine. No need to do anything heavy-weight for now.

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;
+      linearStepVars, reductionVars, alignedVars;
   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);
step[it] =
         firOpBuilder.createConvert(currentLocation, loopVarType, step[it]);
   }
 
   // 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,
-        ifClauseOperand, simdlenClauseOperand, safelenClauseOperand,
+        currentLocation, resultType, lowerBound, upperBound, step, alignedVars,
+        nullptr, ifClauseOperand, simdlenClauseOperand, safelenClauseOperand,
         /*inclusive=*/firOpBuilder.getUnitAttr());
     createBodyOfOp<omp::SimdLoopOp>(SimdLoopOp, converter, currentLocation,
                                     eval, &loopOpClauseList, iv);
     return;
   }
 
   // FIXME: Add support for following clauses:
   // 1. linear

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

let description = [{
     then the upper bound is also included.
 
     The body region can contain any number of blocks. The region is terminated
     by "omp.yield" instruction without operands.
 
     Collapsed loops are represented by the simd-loop having a list of indices,
     bounds and steps where the size of the list is equal to the collapse value.
 
+    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.
 
     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
     SIMD chunk can have a distance in the logical iteration space that is
     greater than or equal to the value given in the clause.
     ```
     omp.simdloop <clauses>
     for (%i1, %i2) : index = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
       // block operations
       omp.yield
     }
     ```
   }];
 
   // 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,
              ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$simdlen,
              ConfinedAttr<OptionalAttr<I64Attr>, [IntPositive]>:$safelen,
              UnitAttr:$inclusive
      );
 
   let regions = (region AnyRegion:$region);
   let assemblyFormat = [{
-    oilist(`if` `(` $if_expr `)`
+    oilist(`aligned` `(`
+              custom<AlignedClause>($aligned_vars, type($aligned_vars),
+                                   $alignment_values) `)`
+          |`if` `(` $if_expr `)`
           |`simdlen` `(` $simdlen  `)`
           |`safelen` `(` $safelen  `)`
     ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
                                   type($step), $inclusive) attr-dict
   }];
 
   let extraClassDeclaration = [{
     /// Returns the number of loops in the simd loop nest.

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

 /// linear-list := linear-val | linear-val linear-list
 /// linear-val := ssa-id-and-type `=` ssa-id-and-type
 static ParseResult
 parseLinearClause(OpAsmParser &parser,
                   SmallVectorImpl<OpAsmParser::UnresolvedOperand> &vars,
                   SmallVectorImpl<Type> &types,
                   SmallVectorImpl<OpAsmParser::UnresolvedOperand> &stepVars) {
   return parser.parseCommaSeparatedList([&]() {
     OpAsmParser::UnresolvedOperand var;
     Type type;
     OpAsmParser::UnresolvedOperand stepVar;
     if (parser.parseOperand(var) || parser.parseEqual() ||
         parser.parseOperand(stepVar) || parser.parseColonType(type))
       return failure();
 
     vars.push_back(var);
     types.push_back(type);
for (unsigned i = 0; i < linearVarsSize; ++i) {
     p << linearVars[i];
     if (linearStepVars.size() > i)
       p << " = " << linearStepVars[i];
     p << " : " << linearVars[i].getType() << separator;
   }
 }
 
 //===----------------------------------------------------------------------===//
+// Parser, verifier and printer for Aligned Clause
+//===----------------------------------------------------------------------===//
+static LogicalResult verifyAlignedClause(Operation *op,
+                                         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())
+      return op->emitOpError()
+             << "expected as many alignment values as aligned variables";
+  } else {
+    if (alignmentValues)
+      return op->emitOpError() << "unexpected alignment values attribute";
+    return success();
+  }
+
+  // Check if each var is aligned only once - OpenMP 4.5 -> 2.8.1 section
+  DenseSet<Value> alignedItems;
+  for (auto it : alignedVariables)
+    if (!alignedItems.insert(it).second)
+      return op->emitOpError() << "aligned variable used more than once";
+
+  if (!alignmentValues)
+    return success();
+
+  // Check if all alignment values are positive - OpenMP 4.5 -> 2.8.1 section
+  for (unsigned i = 0; i < (*alignmentValues).size(); ++i) {
+    if (auto intAttr = (*alignmentValues)[i].dyn_cast<IntegerAttr>()) {
+      if (intAttr.getValue().sle(0))
+        return op->emitOpError() << "alignment should be greater than 0";
+    } else {
+      return op->emitOpError() << "expected integer alignment";
+    }
+  }
+
+  return success();
+}
+
+/// aligned ::= `aligned` `(` aligned-list `)`
+/// aligned-list := aligned-val | aligned-val aligned-list
+/// aligned-val := ssa-id-and-type `->` alignment
+static ParseResult parseAlignedClause(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &alignedItems,
+    SmallVectorImpl<Type> &types, ArrayAttr &alignmentValues) {
+  SmallVector<Attribute> alignmentVec;
+  if (failed(parser.parseCommaSeparatedList([&]() {
+        if (parser.parseOperand(alignedItems.emplace_back()) ||
+            parser.parseColonType(types.emplace_back()) ||
+            parser.parseArrow() ||
+            parser.parseAttribute(alignmentVec.emplace_back())) {
+          return failure();
+        }
+        return success();
+      })))
+    return failure();
+  SmallVector<Attribute> alignments(alignmentVec.begin(), alignmentVec.end());
+  alignmentValues = ArrayAttr::get(parser.getContext(), alignments);
+  return success();
+}
+
+/// Print Aligned Clause
+static void printAlignedClause(OpAsmPrinter &p, Operation *op,
+                               ValueRange alignedVars,
+                               TypeRange alignedVarTypes,
+                               Optional<ArrayAttr> alignmentValues) {
+  for (unsigned i = 0; i < alignedVars.size(); ++i) {
+    if (i != 0)
+      p << ", ";
+    p << alignedVars[i] << " : " << alignedVars[i].getType();
+    p << " -> " << (*alignmentValues)[i];
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // Parser, printer and verifier for Schedule Clause
 //===----------------------------------------------------------------------===//
 
 static ParseResult
 verifyScheduleModifiers(OpAsmParser &parser,
                         SmallVectorImpl<SmallString<12>> &modifiers) {
   if (modifiers.size() > 2)
     return parser.emitError(parser.getNameLoc()) << " unexpected modifier(s)";
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 success();
+  return verifyAlignedClause(*this, this->getAlignmentValues(),
+                             this->getAlignedVars());
 }
 
 //===----------------------------------------------------------------------===//
 // 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>} :
+  }) {operand_segment_sizes = array<i32: 1,1,1,0,0>} :
     (index, index, i32) -> ()
 
   return
 }
 
 // -----
 
+func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index,
+                                       %data_var : memref<i32>) -> () {
+  //  expected-error @below {{expected '->'}}
+  omp.simdloop aligned(%data_var : memref<i32>)
+  for (%iv) : index = (%lb) to (%ub) step (%step) {
+    omp.yield
+  }
+  return
+}
+
+// -----
+
+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>) -> ()
+  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>) -> ()
+  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) -> ()
+  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>) -> ()
+  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>) -> ()
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_pretty_simdlen(%lb : index, %ub : index, %step : index) -> () {
   // expected-error @below {{op attribute 'simdlen' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive}}
   omp.simdloop simdlen(0) for (%iv): index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
   return
 }
 

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>} :
+  }) {operand_segment_sizes = array<i32: 1,1,1,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>) -> ()
+  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>) -> ()
+  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
   }
   return
 }
 
+// CHECK-LABEL:   func.func @omp_simdloop_pretty_aligned(
+func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index,
+                                       %data_var : memref<i32>,
+                                       %data_var1 : memref<i32>) -> () {
+  // CHECK:      omp.simdloop   aligned(%{{.*}} : memref<i32> -> 32 : i64,
+  // CHECK-SAME: %{{.*}} : memref<i32> -> 128 : i64)
+  // CHECK-SAME: for  (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
+  omp.simdloop aligned(%data_var :  memref<i32> -> 32, %data_var1 : memref<i32> -> 128)
+    for (%iv) : index = (%lb) to (%ub) step (%step) {
+      omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_pretty_if
 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
 }

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>} :
+  }) {operand_segment_sizes = array<i32: 1,1,1,0,0>} :
     (i64, i64, i64) -> ()
 
   llvm.return
 }
 // CHECK: llvm.loop.parallel_accesses
 // CHECK-NEXT: llvm.loop.vectorize.enable
 
 // -----