[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?

416

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

417

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

138

OpAsmParser has a parseCommaSeparatedList, can that be used instead?

mlir/test/Dialect/OpenMP/ops.mlir

356

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

417

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

138

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

216

Nit:

223

Nit: Can you add a test for this?

231

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.

242

Nit: Can you add a test for this?

[domada]

Applied remarks. Added more tests.

[domada]

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

223

Added test case: omp_simdloop_unexpected_alignment

231

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.

242

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

231

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, alignmentVars;
   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,
+        alignmentVars, 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_val` operand additionally specifies alignment of each
+    corresponding aligned operand. Note that `$aligned_vars` and
+    `alignment_val` variadic lists 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<AnyType>:$aligned_vars,
+             Variadic<I32>:$alignment_val,
              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(`if` `(` $if_expr `)` |
+             `aligned` `(`
+              custom<AlignedClause>($aligned_vars, type($aligned_vars),
+                                   $alignment_val) `)`
           |`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

 }
 
 template <typename ClauseAttr>
 void printClauseAttr(OpAsmPrinter &p, Operation *op, ClauseAttr attr) {
   p << stringifyEnum(attr.getValue());
 }
 
 //===----------------------------------------------------------------------===//
-// Parser and printer for Linear Clause
+// Parser and printer for Linear or Aligned Clauses comma list
 //===----------------------------------------------------------------------===//
 
-/// linear ::= `linear` `(` linear-list `)`
-/// 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) {
+/// comma-list := comma-val | comma-val comma-list
+/// comma-val := ssa-id-and-type `=` ssa-id-and-type
+static ParseResult parseCommaSeparatedListClause(
+    OpAsmParser &parser, SmallVectorImpl<OpAsmParser::UnresolvedOperand> &items,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &itemsVars) {
   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);
+    items.push_back(var);
     types.push_back(type);
-    stepVars.push_back(stepVar);
+    itemsVars.push_back(stepVar);
     return success();
   });
 }
 
-/// Print Linear Clause
+static void printCommaSeparatedListClause(OpAsmPrinter &p, Operation *op,
+                                          ValueRange items, TypeRange types,
+                                          ValueRange itemsVars) {
+  size_t itemsSize = items.size();
+  for (unsigned i = 0; i < itemsSize; ++i) {
+    std::string separator = i == itemsSize - 1 ? "" : ", ";
+    p << items[i];
+    if (itemsVars.size() > i)
+      p << " = " << itemsVars[i];
+    p << " : " << items[i].getType() << separator;
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// Parser and printer for Aligned Clause
+//===----------------------------------------------------------------------===//
+/// linear ::= `aligned` `(` comma-list `)`
+static ParseResult parseAlignedClause(
+    OpAsmParser &parser,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &alignedItems,
+    SmallVectorImpl<Type> &types,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &alignmentValues) {
+  return parseCommaSeparatedListClause(parser, alignedItems, types,
+                                       alignmentValues);
+}
+
+static void printAlignedClause(OpAsmPrinter &p, Operation *op,
+                               ValueRange alignedVars,
+                               TypeRange alignedVarTypes,
+                               ValueRange alignmentValues) {
+  printCommaSeparatedListClause(p, op, alignedVars, alignedVarTypes,
+                                alignmentValues);
+}
+
+//===----------------------------------------------------------------------===//
+// Parser and printer for Linear Clause
+//===----------------------------------------------------------------------===//
+/// linear ::= `linear` `(` comma-list `)`
+static ParseResult
+parseLinearClause(OpAsmParser &parser,
+                  SmallVectorImpl<OpAsmParser::UnresolvedOperand> &vars,
+                  SmallVectorImpl<Type> &types,
+                  SmallVectorImpl<OpAsmParser::UnresolvedOperand> &stepVars) {
+  return parseCommaSeparatedListClause(parser, vars, types, stepVars);
+}
+
 static void printLinearClause(OpAsmPrinter &p, Operation *op,
                               ValueRange linearVars, TypeRange linearVarTypes,
                               ValueRange linearStepVars) {
-  size_t linearVarsSize = linearVars.size();
-  for (unsigned i = 0; i < linearVarsSize; ++i) {
-    std::string separator = i == linearVarsSize - 1 ? "" : ", ";
-    p << linearVars[i];
-    if (linearStepVars.size() > i)
-      p << " = " << linearStepVars[i];
-    p << " : " << linearVars[i].getType() << separator;
-  }
+  printCommaSeparatedListClause(p, op, linearVars, linearVarTypes,
+                                linearStepVars);
 }
 
 //===----------------------------------------------------------------------===//
 // 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)";
   for (const auto &mod : modifiers) {
     // Translate the string. If it has no value, then it was not a valid
     // modifier!
     auto symbol = symbolizeScheduleModifier(mod);
     if (!symbol)
       return parser.emitError(parser.getNameLoc())
              << " unknown modifier type: " << mod;
   }
 
   // If we have one modifier that is "simd", then stick a "none" modiifer in
   // index 0.
   if (modifiers.size() == 1) {
     if (symbolizeScheduleModifier(modifiers[0]) == ScheduleModifier::simd) {
       modifiers.push_back(modifiers[0]);
       modifiers[0] = stringifyScheduleModifier(ScheduleModifier::none);
     }
   } else if (modifiers.size() == 2) {
     // If there are two modifier:
     // First modifier should not be simd, second one should be simd
     if (symbolizeScheduleModifier(modifiers[0]) == ScheduleModifier::simd ||
         symbolizeScheduleModifier(modifiers[1]) != ScheduleModifier::simd)
       return parser.emitError(parser.getNameLoc())
              << " incorrect modifier order";
   }
   return success();
 }
 
 /// schedule ::= `schedule` `(` sched-list `)`
 /// sched-list ::= sched-val | sched-val sched-list |
 ///                sched-val `,` sched-modifier
 /// sched-val ::= sched-with-chunk | sched-wo-chunk
 /// sched-with-chunk ::= sched-with-chunk-types (`=` ssa-id-and-type)?

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,0>} :
     (index, index, i32) -> ()
 
   return
 }
 
 // -----
+func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %alignment_var : 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_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,0>} :
     (index, index, index) -> ()
 
   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: omp_simdloop_pretty_aligned
+func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %alignment_var : i32) -> () {
+
+  // CHECK: omp.simdloop aligned(%{{.*}} = %{{.*}} : memref<i32>)
+  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.simdloop aligned(%data_var = %alignment_var : memref<i32>)
+  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,0>} :
     (i64, i64, i64) -> ()
 
   llvm.return
 }
 // CHECK: llvm.loop.parallel_accesses
 // CHECK-NEXT: llvm.loop.vectorize.enable
 
 // -----