[Flang][MLIR][OpenMP] Add support for simdlen clause

[psoni2628]

Summary

This supports lowering from parse-tree to MLIR and translation from
MLIR to LLVM IR using OMPIRBuilder for OpenMP simdlen clause in SIMD
construct.

[arnamoy10]

Thanks for working on this patch. Can we add a verifier for
simdlen to error when the given value is negative? From the OpenMP standards "The parameter of the simdlen clause must be a constant positive integer expression."

[psoni2628]

In D130195#3666505, @arnamoy10 wrote:

Thanks for working on this patch. Can we add a verifier for
simdlen to error when the given value is negative? From the OpenMP standards "The parameter of the simdlen clause must be a constant positive integer expression."

Good point. I think it would be simpler to change simdlen argument type from I32 to Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]> in the MLIR op definition. We use the aforementioned type currently for clauses such as ordered, which are also required to be "constant positive integer expression".

[psoni2628]

• Use IntPositive to ensure that simdlen is a constant, positive integer

[arnamoy10]

1. Can we have a test case in invalid.mlir
2. Can we have a testcase where the simdlen is specified through a function argument or a local variable?

[psoni2628]

• Add invalid.mlir testcase
• Add Flang test case with simdlen from a local parameter instead of a literal

[arnamoy10]

LGTM.

[psoni2628]

• Rebase so that the patch can be applied

[shraiysh]

flang/lib/Lower/OpenMP.cpp

907–909

1049–1050

We can have simdlenClauseOperand as an IntegerAttr instead of an Attribute (suggested edit in the code above). That will avoid type-casting to and from Attribute for this.

flang/test/Lower/OpenMP/simd.f90

61

Can we also have a test here that has an expression. Like SIMDLEN(simdlen * 2 + n) or something.

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

975

nit: expand auto.

978

There are no tests for MLIR->LLVM IR translation. Changes to OpenMPToLLVMIRTranslation.cpp and the tests for that (in mlir/test/Target/LLVMIR/openmp-llvm.mlir) should be a part of the same patch. You can send them as a separate patch if you'd like.

[psoni2628]

• Expand auto
• Add expression Flang test case
• Use IntegerAttr directly to avoid the need for a dyn_cast
• Add mlir/test/Target/LLVMIR/openmp-llvm.mlir testcase

flang/test/Lower/OpenMP/simd.f90

61

I instead used SIMDLEN(simdlen*2 + 2), because in this example n is not constant, so you can't use it in an expression for simdlen.

[shraiysh]

LGTM.

flang/test/Lower/OpenMP/simd.f90

61

Hmm, gfortran seems to accept simdlen*2+n in this example. The standard mentions that the argument must be a "constant positive integer expression". It is probably incorrect behavior in gfortran. Please wait for a day or two if anyone has a preference for being accurate to standard vs being compatible with gfortran. If no one objects, we can go ahead and merge this. For me, this is fine.

[psoni2628]

flang/test/Lower/OpenMP/simd.f90

61

I wanted to point out that it is the parser that complains with Must be a constant value. If we are going to try and be compatible with gfortran, I think it would make more sense to do this in a separate patch where we change the parser and the other affected parts accordingly.

Also, if you write a similar C program as shown below and compile it with Clang, Clang also complains similarly, so I believe Flang's existing parsing behaviour is correct.

#define N  1000000000

void func (int n) {

	int a[N] = {0};

	#pragma omp simd safelen(n)
	for (int i=0; i<N; i++) {
		a[i] = a[i] + i;
	}
}

./ex.c:10:27: error: expression is not an integer constant expression
        #pragma omp simd safelen(n)

[peixin]

Nit: In summary:

This supports lowering from parse-tree to MLIR and translation from MLIR to LLVM IR using OMPIRBuilder for OpenMP simdlen clause in SIMD construct.

flang/test/Lower/OpenMP/simd.f90

61

@shraiysh @psoni2628 Don't worry about this. It's safe to emit one semantic error for this. Check the following cases:

subroutine sub(N, x, y, z)
  integer :: N
  integer :: l = 1
  integer :: x(N), y(N), z(N)
  !$omp simd simdlen(l)
  do i = 1, N
    x(i) = x(i) + y(i)*z(i)
  enddo
  !$omp end simd
end

subroutine sub(N, x, y, z)
  integer :: N
  integer, parameter :: l = 1
  integer :: x(N), y(N), z(N)
  !$omp simd simdlen(l)
  do i = 1, N
    x(i) = x(i) + y(i)*z(i)
  enddo
  !$omp end simd
end

Run gfortran -fopenmp test.f90 -O3 -S -o test-omp.s and you will see simdlen with variable won't work.

The variable is unknown in compile-time. I don't think you can add one variable in the loop vectorization metadata.

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;
   mlir::Value scheduleChunkClauseOperand, ifClauseOperand;
   mlir::Attribute scheduleClauseOperand, noWaitClauseOperand,
       orderedClauseOperand, orderClauseOperand;
+  mlir::IntegerAttr simdlenClauseOperand;
   SmallVector<Attribute> reductionDeclSymbols;
   Fortran::lower::StatementContext stmtCtx;
   const auto &loopOpClauseList = std::get<Fortran::parser::OmpClauseList>(
       std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t);
 
   const auto ompDirective =
       std::get<Fortran::parser::OmpLoopDirective>(
           std::get<Fortran::parser::OmpBeginLoopDirective>(loopConstruct.t).t)
if (const auto &scheduleClause =
                   firOpBuilder.getContext(), decl.sym_name()));
             }
           }
         }
       } else {
         TODO(currentLocation,
              "Reduction of intrinsic procedures is not supported");
       }
+    } else if (const auto &simdlenClause =
+                   std::get_if<Fortran::parser::OmpClause::Simdlen>(
+                       &clause.u)) {
+      const auto *expr = Fortran::semantics::GetExpr(simdlenClause->v);
+      const std::optional<std::int64_t> simdlenVal =
+          Fortran::evaluate::ToInt64(*expr);
+      simdlenClauseOperand = firOpBuilder.getI64IntegerAttr(*simdlenVal);
     }
   }
 
   // The types of lower bound, upper bound, and step are converted into the
   // type of the loop variable if necessary.
   mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
   for (unsigned it = 0; it < (unsigned)lowerBound.size(); it++) {
     lowerBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 lowerBound[it]);
     upperBound[it] = firOpBuilder.createConvert(currentLocation, loopVarType,
                                                 upperBound[it]);
     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, /*inclusive=*/firOpBuilder.getUnitAttr());
+        ifClauseOperand, simdlenClauseOperand,
+        /*inclusive=*/firOpBuilder.getUnitAttr());
     createBodyOfOp<omp::SimdLoopOp>(SimdLoopOp, converter, currentLocation,
                                     eval, &loopOpClauseList, iv);
     return;
   }
 
   // FIXME: Add support for following clauses:
   // 1. linear
   // 2. order

flang/test/Lower/OpenMP/simd.f90

integer :: i, n, threshold
   do i = 1, n
     ! CHECK: fir.store %[[I]] to %[[LOCAL:.*]] : !fir.ref<i32>
     ! CHECK: %[[LD:.*]] = fir.load %[[LOCAL]] : !fir.ref<i32>
     ! CHECK: fir.call @_FortranAioOutputInteger32({{.*}}, %[[LD]]) : (!fir.ref<i8>, i32) -> i1
     print*, i
   end do
   !$OMP END SIMD
 end subroutine
+
+!CHECK-LABEL: func @_QPsimdloop_with_simdlen_clause
+subroutine simdloop_with_simdlen_clause(n, threshold)
+integer :: i, n, threshold
+  !$OMP SIMD SIMDLEN(2)
+  ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
+  ! CHECK: %[[UB:.*]] = fir.load %arg0
+  ! CHECK: %[[STEP:.*]] = arith.constant 1 : i32
+  ! CHECK: omp.simdloop simdlen(2) for (%[[I:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) inclusive  step (%[[STEP]]) {
+  do i = 1, n
+    ! CHECK: fir.store %[[I]] to %[[LOCAL:.*]] : !fir.ref<i32>
+    ! CHECK: %[[LD:.*]] = fir.load %[[LOCAL]] : !fir.ref<i32>
+    ! CHECK: fir.call @_FortranAioOutputInteger32({{.*}}, %[[LD]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  !$OMP END SIMD
+end subroutine
+
+!CHECK-LABEL: func @_QPsimdloop_with_simdlen_clause_from_param
+subroutine simdloop_with_simdlen_clause_from_param(n, threshold)
+integer :: i, n, threshold
+integer, parameter :: simdlen = 2;
+  !$OMP SIMD SIMDLEN(simdlen)
+  ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
+  ! CHECK: %[[UB:.*]] = fir.load %arg0
+  ! CHECK: %[[STEP:.*]] = arith.constant 1 : i32
+  ! CHECK: omp.simdloop simdlen(2) for (%[[I:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) inclusive  step (%[[STEP]]) {
+  do i = 1, n
+    ! CHECK: fir.store %[[I]] to %[[LOCAL:.*]] : !fir.ref<i32>
+    ! CHECK: %[[LD:.*]] = fir.load %[[LOCAL]] : !fir.ref<i32>
+    ! CHECK: fir.call @_FortranAioOutputInteger32({{.*}}, %[[LD]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  !$OMP END SIMD
+end subroutine
+
+!CHECK-LABEL: func @_QPsimdloop_with_simdlen_clause_from_expr_from_param
+subroutine simdloop_with_simdlen_clause_from_expr_from_param(n, threshold)
+integer :: i, n, threshold
+integer, parameter :: simdlen = 2;
+  !$OMP SIMD SIMDLEN(simdlen*2 + 2)
+  ! CHECK: %[[LB:.*]] = arith.constant 1 : i32
+  ! CHECK: %[[UB:.*]] = fir.load %arg0
+  ! CHECK: %[[STEP:.*]] = arith.constant 1 : i32
+  ! CHECK: omp.simdloop simdlen(6) for (%[[I:.*]]) : i32 = (%[[LB]]) to (%[[UB]]) inclusive  step (%[[STEP]]) {
+  do i = 1, n
+    ! CHECK: fir.store %[[I]] to %[[LOCAL:.*]] : !fir.ref<i32>
+    ! CHECK: %[[LD:.*]] = fir.load %[[LOCAL]] : !fir.ref<i32>
+    ! CHECK: fir.call @_FortranAioOutputInteger32({{.*}}, %[[LD]]) : (!fir.ref<i8>, i32) -> i1
+    print*, i
+  end do
+  !$OMP END SIMD
+end subroutine

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.
 
     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.
     ```
     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,
              Optional<I1>:$if_expr,
+             Confined<OptionalAttr<I64Attr>, [IntPositive]>:$simdlen,
              UnitAttr:$inclusive
      );
  
   let regions = (region AnyRegion:$region);
   let assemblyFormat = [{
     oilist(`if` `(` $if_expr `)`
+          |`simdlen` `(` $simdlen  `)`
     ) `for` custom<LoopControl>($region, $lowerBound, $upperBound, $step,
                                   type($step), $inclusive) attr-dict
   }];
 
   let extraClassDeclaration = [{
     /// Returns the number of loops in the simd loop nest.
     unsigned getNumLoops() { return lowerBound().size(); }
 

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

if (failed(bodyGenStatus))
       return failure();
   }
 
   // Collapse loops.
   llvm::IRBuilderBase::InsertPoint afterIP = loopInfos.front()->getAfterIP();
   llvm::CanonicalLoopInfo *loopInfo =
       ompBuilder->collapseLoops(ompLoc.DL, loopInfos, {});
 
-  ompBuilder->applySimd(loopInfo, nullptr);
+  llvm::ConstantInt *simdlen = nullptr;
+  if (llvm::Optional<uint64_t> simdlenVar = loop.simdlen())
+    simdlen = builder.getInt64(simdlenVar.getValue());
+
+  ompBuilder->applySimd(loopInfo, simdlen);
 
   builder.restoreIP(afterIP);
   return success();
 }
 
 /// Convert an Atomic Ordering attribute to llvm::AtomicOrdering.
 llvm::AtomicOrdering
 convertAtomicOrdering(Optional<omp::ClauseMemoryOrderKind> ao) {

mlir/test/Dialect/OpenMP/invalid.mlir

func.func @omp_simdloop(%lb : index, %ub : index, %step : i32) -> () {
   }) {operand_segment_sizes = dense<[1,1,1,0]> : vector<4xi32>} :
     (index, index, 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
+}
+
+// -----
+
 // expected-error @below {{op expects initializer region with one argument of the reduction type}}
 omp.reduction.declare @add_f32 : f64
 init {
 ^bb0(%arg: f32):
   %0 = arith.constant 0.0 : f32
   omp.yield (%0 : f32)
 }
 combiner {

mlir/test/Dialect/OpenMP/ops.mlir

 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: omp_simdloop_pretty_simdlen
+func.func @omp_simdloop_pretty_simdlen(%lb : index, %ub : index, %step : index) -> () {
+  // CHECK: omp.simdloop simdlen(2) for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  omp.simdloop simdlen(2) for (%iv): index = (%lb) to (%ub) step (%step) {
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_pretty_multiple
 func.func @omp_simdloop_pretty_multiple(%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> () {
   // CHECK: omp.simdloop for (%{{.*}}, %{{.*}}) : index = (%{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}})
   omp.simdloop for (%iv1, %iv2) : index = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
     omp.yield
   }
   return
 }

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

llvm.func @simdloop_simple_multiple(%lb1 : i64, %ub1 : i64, %step1 : i64, %lb2 : i64, %ub2 : i64, %step2 : i64, %arg0: !llvm.ptr<f32>, %arg1: !llvm.ptr<f32>) {
   } 
   llvm.return
 }
 // CHECK: llvm.loop.parallel_accesses
 // CHECK-NEXT: llvm.loop.vectorize.enable
 
 // -----
 
+// CHECK-LABEL: @simdloop_simple_multiple_simdlen
+llvm.func @simdloop_simple_multiple_simdlen(%lb1 : i64, %ub1 : i64, %step1 : i64, %lb2 : i64, %ub2 : i64, %step2 : i64, %arg0: !llvm.ptr<f32>, %arg1: !llvm.ptr<f32>) {
+  omp.simdloop simdlen(2) for (%iv1, %iv2) : i64 = (%lb1, %lb2) to (%ub1, %ub2) step (%step1, %step2) {
+    %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
+    // CHECK-NEXT: llvm.access.group
+    %4 = llvm.getelementptr %arg0[%iv1] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
+    %5 = llvm.getelementptr %arg1[%iv2] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
+    llvm.store %3, %4 : !llvm.ptr<f32>
+    llvm.store %3, %5 : !llvm.ptr<f32>
+    omp.yield
+  }
+  llvm.return
+}
+// CHECK: llvm.loop.parallel_accesses
+// CHECK-NEXT: llvm.loop.vectorize.enable
+// CHECK-NEXT: llvm.loop.vectorize.width{{.*}}i64 2
+
+// -----
+
 llvm.func @body(i64)
 
 llvm.func @test_omp_wsloop_ordered(%lb : i64, %ub : i64, %step : i64) -> () {
  omp.wsloop ordered(0)
  for (%iv) : i64 = (%lb) to (%ub) step (%step) {
   // CHECK: call void @__kmpc_dispatch_init_8u(ptr @{{.*}}, i32 %{{.*}}, i32 66, i64 1, i64 %{{.*}}, i64 1, i64 1)
   // CHECK: call void @__kmpc_dispatch_fini_8u
   // CHECK: %[[continue:.*]] = call i32 @__kmpc_dispatch_next_8u