[NFC][OpenMP] Fix worksharing-loop

[peixin]

Summary

1. Remove the redundant collapse clause in MLIR OpenMP worksharing-loop operation.
2. Fix several typos.
3. Refactor the chunk size type conversion since CreateSExtOrTrunc has both type check and type conversion.

[kiranchandramohan]

Have we handled the ordered clause with n>0? If the representation for that is similar to collapse then we might need some way to distinguish collapse and ordered loops.

[peixin]

In D128338#3601472, @kiranchandramohan wrote:

Have we handled the ordered clause with n>0? If the representation for that is similar to collapse then we might need some way to distinguish collapse and ordered loops.

Not yet. I had one try (https://reviews.llvm.org/D116292#change-bQfuPOutx59a), but there seems to be some problem with overflow. @Meinersbur is working on that. But I assume the implementation will be different with collapse clause.

[kiranchandramohan]

In D128338#3601529, @peixin wrote:

In D128338#3601472, @kiranchandramohan wrote:

Have we handled the ordered clause with n>0? If the representation for that is similar to collapse then we might need some way to distinguish collapse and ordered loops.

Not yet. I had one try (https://reviews.llvm.org/D116292#change-bQfuPOutx59a), but there seems to be some problem with overflow. @Meinersbur is working on that. But I assume the implementation will be different with collapse clause.

Maybe I should have mentioned this in D116292, the standard specifies that for ordered(n), there are n loops that are associated with the ordered clause. One way to represent this in the IR is like the collapsed loop representation. If we don't bunch up the loops into a single operation then there is a possibility that FIR or other transformations can modify one of these loops which should not be permitted. So we might end up with a similar representation for collapsed and ordered loops, then we need to distinguish which one is collapse and which one is ordered. This can be by attributes, enums, or even additional information in the representation of the ordered clause (whose absence makes it a collapsed loop). Alternatively, we can leave this for now until we decide on the ordered loop representation.

[peixin]

One way to represent this in the IR is like the collapsed loop representation. If we don't bunch up the loops into a single operation then there is a possibility that FIR or other transformations can modify one of these loops which should not be permitted. So we might end up with a similar representation for collapsed and ordered loops, then we need to distinguish which one is collapse and which one is ordered. This can be by attributes, enums, or even additional information in the representation of the ordered clause (whose absence makes it a collapsed loop).

Agree. Michael mentioned there should be some transformation in canonical loop for the doacross loop nest. I am not sure if he plans to do it before MLIR or in OMPIRBuilder, so I abandon the whole patch including the doacross loops info in WsLoop operation. In addition, the ordered depend clause should be adjusted if necessary when supporting doacross loop nest for ordered clause (with n). This is why I also abandon the lowering of ordered depend directive in fir-dev. I have mentioned this to Michael and he took up both.

Alternatively, we can leave this for now until we decide on the ordered loop representation.

According to the above, I think we should leave this for now since we cannot have a final decision for the ordered loop representation if not supporting the complete doacross nest (including ordered clause and ordered depend) for now.

[peixin]

@kiranchandramohan Is it OK to land this now?

[kiranchandramohan]

LGTM. One Nit comment.

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

328

Nit: Can you instead add the following?

Collapsed loops are represented by the worksharing loop having a list of indices, bounds and steps where the size of the list is equal to the collapse value.

[peixin]

Add the description for collapsed loops as @kiranchandramohan suggested.

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;
-  mlir::Attribute scheduleClauseOperand, collapseClauseOperand,
-      noWaitClauseOperand, orderedClauseOperand, orderClauseOperand;
+  mlir::Attribute scheduleClauseOperand, noWaitClauseOperand,
+      orderedClauseOperand, orderClauseOperand;
   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)
           .v;
   if (llvm::omp::OMPD_parallel_do == ompDirective) {
static void genOMP(Fortran::lower::AbstractConverter &converter,
   // 1. linear
   // 2. order
   auto wsLoopOp = firOpBuilder.create<mlir::omp::WsLoopOp>(
       currentLocation, lowerBound, upperBound, step, linearVars, linearStepVars,
       reductionVars, /*reductions=*/nullptr,
       scheduleClauseOperand.dyn_cast_or_null<omp::ClauseScheduleKindAttr>(),
       scheduleChunkClauseOperand, /*schedule_modifiers=*/nullptr,
       /*simd_modifier=*/nullptr,
-      collapseClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       noWaitClauseOperand.dyn_cast_or_null<UnitAttr>(),
       orderedClauseOperand.dyn_cast_or_null<IntegerAttr>(),
       orderClauseOperand.dyn_cast_or_null<omp::ClauseOrderKindAttr>(),
       /*inclusive=*/firOpBuilder.getUnitAttr());
 
   // Handle attribute based clauses.
   for (const Fortran::parser::OmpClause &clause : loopOpClauseList.v) {
     if (const auto &orderedClause =
             std::get_if<Fortran::parser::OmpClause::Ordered>(&clause.u)) {
       if (orderedClause->v.has_value()) {
         const auto *expr = Fortran::semantics::GetExpr(orderedClause->v);
         const std::optional<std::int64_t> orderedClauseValue =
             Fortran::evaluate::ToInt64(*expr);
         wsLoopOp.ordered_valAttr(
             firOpBuilder.getI64IntegerAttr(*orderedClauseValue));
       } else {
         wsLoopOp.ordered_valAttr(firOpBuilder.getI64IntegerAttr(0));
       }
-    } else if (const auto &collapseClause =
-                   std::get_if<Fortran::parser::OmpClause::Collapse>(
-                       &clause.u)) {
-      const auto *expr = Fortran::semantics::GetExpr(collapseClause->v);
-      const std::optional<std::int64_t> collapseValue =
-          Fortran::evaluate::ToInt64(*expr);
-      wsLoopOp.collapse_valAttr(firOpBuilder.getI64IntegerAttr(*collapseValue));
     } else if (const auto &scheduleClause =
                    std::get_if<Fortran::parser::OmpClause::Schedule>(
                        &clause.u)) {
       mlir::MLIRContext *context = firOpBuilder.getContext();
       const auto &scheduleType = scheduleClause->v;
       const auto &scheduleKind =
           std::get<Fortran::parser::OmpScheduleClause::ScheduleType>(
               scheduleType.t);

flang/test/Lower/OpenMP/omp-wsloop-collapse.f90

 ! CHECK:           %[[VAL_24:.*]] = fir.load %[[VAL_1]] : !fir.ref<i32>
 ! CHECK:           %[[VAL_25:.*]] = arith.constant 1 : i32
 ! CHECK:           %[[VAL_26:.*]] = arith.constant 1 : i32
 ! CHECK:           %[[VAL_27:.*]] = fir.load %[[VAL_2]] : !fir.ref<i32>
 ! CHECK:           %[[VAL_28:.*]] = arith.constant 1 : i32
   do i = 1, a
      do j= 1, b
         do k = 1, c
-! CHECK:           omp.wsloop collapse(3) for (%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]]) : i32 = (%[[VAL_20]], %[[VAL_23]], %[[VAL_26]]) to (%[[VAL_21]], %[[VAL_24]], %[[VAL_27]]) inclusive step (%[[VAL_22]], %[[VAL_25]], %[[VAL_28]]) {
+! CHECK:           omp.wsloop for (%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]]) : i32 = (%[[VAL_20]], %[[VAL_23]], %[[VAL_26]]) to (%[[VAL_21]], %[[VAL_24]], %[[VAL_27]]) inclusive step (%[[VAL_22]], %[[VAL_25]], %[[VAL_28]]) {
 ! CHECK:             fir.store %[[ARG0]] to %[[STORE_IV0:.*]] : !fir.ref<i32>
 ! CHECK:             fir.store %[[ARG1]] to %[[STORE_IV1:.*]] : !fir.ref<i32>
 ! CHECK:             fir.store %[[ARG2]] to %[[STORE_IV2:.*]] : !fir.ref<i32>
 ! CHECK:             %[[VAL_12:.*]] = fir.load %[[VAL_6]] : !fir.ref<i32>
 ! CHECK:             %[[LOAD_IV0:.*]] = fir.load %[[STORE_IV0]] : !fir.ref<i32>
 ! CHECK:             %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[LOAD_IV0]] : i32
 ! CHECK:             %[[LOAD_IV1:.*]] = fir.load %[[STORE_IV1]] : !fir.ref<i32>
 ! CHECK:             %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[LOAD_IV1]] : i32

flang/test/Lower/OpenMP/omp-wsloop-variable.f90

 !CHECK:  %[[TMP0:.*]] = arith.constant 1 : i32
 !CHECK:  %[[TMP1:.*]] = arith.constant 100 : i32
 !CHECK:  %[[TMP2:.*]] = fir.convert %[[TMP0]] : (i32) -> i64
 !CHECK:  %[[TMP3:.*]] = fir.convert %{{.*}} : (i8) -> i64
 !CHECK:  %[[TMP4:.*]] = fir.convert %{{.*}} : (i16) -> i64
 !CHECK:  %[[TMP5:.*]] = fir.convert %{{.*}} : (i128) -> i64
 !CHECK:  %[[TMP6:.*]] = fir.convert %[[TMP1]] : (i32) -> i64
 !CHECK:  %[[TMP7:.*]] = fir.convert %{{.*}} : (i32) -> i64
-!CHECK:  omp.wsloop collapse(2) for (%[[ARG0:.*]], %[[ARG1:.*]]) : i64 = (%[[TMP2]], %[[TMP5]]) to (%[[TMP3]], %[[TMP6]]) inclusive step (%[[TMP4]], %[[TMP7]]) {
+!CHECK:  omp.wsloop for (%[[ARG0:.*]], %[[ARG1:.*]]) : i64 = (%[[TMP2]], %[[TMP5]]) to (%[[TMP3]], %[[TMP6]]) inclusive step (%[[TMP4]], %[[TMP7]]) {
 !CHECK:    fir.store %[[ARG0]] to %[[STORE_IV0:.*]] : !fir.ref<i64>
 !CHECK:    fir.store %[[ARG1]] to %[[STORE_IV1:.*]] : !fir.ref<i64>
 !CHECK:    %[[LOAD_IV0:.*]] = fir.load %[[STORE_IV0]] : !fir.ref<i64>
 !CHECK:    %[[LOAD_IV1:.*]] = fir.load %[[STORE_IV1]] : !fir.ref<i64>
 !CHECK:    %[[TMP10:.*]] = arith.addi %[[LOAD_IV0]], %[[LOAD_IV1]] : i64
 !CHECK:    %[[TMP11:.*]] = fir.convert %[[TMP10]] : (i64) -> f32
 !CHECK:    fir.store %[[TMP11]] to %{{.*}} : !fir.ref<f32>
 !CHECK:    omp.yield

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

let description = [{
     }
     ```
 
     The `linear_step_vars` operand additionally specifies the step for each
     associated linear operand. Note that the `linear_vars` and
     `linear_step_vars` variadic lists should contain the same number of
     elements.
 
-    Reductions can be performed in a workshare loop by specifying reduction
+    Reductions can be performed in a worksharing-loop by specifying reduction
     accumulator variables in `reduction_vars` and symbols referring to reduction
     declarations in the `reductions` attribute. Each reduction is identified
     by the accumulator it uses and accumulators must not be repeated in the same
     reduction. The `omp.reduction` operation accepts the accumulator and a
     partial value which is considered to be produced by the current loop
     iteration for the given reduction. If multiple values are produced for the
     same accumulator, i.e. there are multiple `omp.reduction`s, the last value
     is taken. The reduction declaration specifies how to combine the values from
     each iteration into the final value, which is available in the accumulator
     after the loop completes.
 
     The optional `schedule_val` attribute specifies the loop schedule for this
     loop, determining how the loop is distributed across the parallel threads.
     The optional `schedule_chunk_var` associated with this determines further
     controls this distribution.
 
-    The optional `collapse_val` attribute specifies the number of loops which
-    are collapsed to form the worksharing loop.
+    Collapsed loops are represented by the worksharing-loop having a list of
+    indices, bounds and steps where the size of the list is equal to the
+    collapse value.
 
     The `nowait` attribute, when present, signifies that there should be no
     implicit barrier at the end of the loop.
 
     The optional `ordered_val` attribute specifies how many loops are associated
     with the worksharing-loop construct. The value of zero refers to the ordered
     clause specified without parameter.
 
     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".
   }];
 
   let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              Variadic<IntLikeType>:$upperBound,
              Variadic<IntLikeType>:$step,
              Variadic<AnyType>:$linear_vars,
              Variadic<I32>:$linear_step_vars,
              Variadic<OpenMP_PointerLikeType>:$reduction_vars,
              OptionalAttr<SymbolRefArrayAttr>:$reductions,
              OptionalAttr<ScheduleKindAttr>:$schedule_val,
              Optional<AnyType>:$schedule_chunk_var,
              OptionalAttr<ScheduleModifierAttr>:$schedule_modifier,
              UnitAttr:$simd_modifier,
-             Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
              UnitAttr:$nowait,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
              OptionalAttr<OrderKindAttr>:$order_val,
              UnitAttr:$inclusive);
 
   let builders = [
     OpBuilder<(ins "ValueRange":$lowerBound, "ValueRange":$upperBound,
                "ValueRange":$step,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
   ];
 
   let regions = (region AnyRegion:$region);
 
   let extraClassDeclaration = [{
-    /// Returns the number of loops in the workshape loop nest.
+    /// Returns the number of loops in the worksharing-loop nest.
     unsigned getNumLoops() { return lowerBound().size(); }
 
     /// Returns the number of reduction variables.
     unsigned getNumReductionVars() { return reduction_vars().size(); }
 
     // TODO: remove this once emitAccessorPrefix is set to
     // kEmitAccessorPrefix_Prefixed for the dialect.
     /// Returns the reduction variables
     operand_range getReductionVars() { return reduction_vars(); }
   }];
   let hasCustomAssemblyFormat = 1;
   let assemblyFormat = [{
     oilist(`linear` `(`
               custom<LinearClause>($linear_vars, type($linear_vars),
                                    $linear_step_vars) `)`
           |`schedule` `(`
               custom<ScheduleClause>(
                 $schedule_val, $schedule_modifier, $simd_modifier,
                 $schedule_chunk_var, type($schedule_chunk_var)) `)`
-          |`collapse` `(` $collapse_val `)`
           |`nowait` $nowait
           |`ordered` `(` $ordered_val `)`
           |`order` `(` custom<ClauseAttr>($order_val) `)`
           |`reduction` `(`
               custom<ReductionVarList>(
                 $reduction_vars, type($reduction_vars), $reductions
               ) `)`
     ) `for` custom<WsLoopControl>($region, $lowerBound, $upperBound, $step,

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

 void WsLoopOp::build(OpBuilder &builder, OperationState &state,
                      ValueRange lowerBound, ValueRange upperBound,
                      ValueRange step, ArrayRef<NamedAttribute> attributes) {
   build(builder, state, lowerBound, upperBound, step,
         /*linear_vars=*/ValueRange(),
         /*linear_step_vars=*/ValueRange(), /*reduction_vars=*/ValueRange(),
         /*reductions=*/nullptr, /*schedule_val=*/nullptr,
         /*schedule_chunk_var=*/nullptr, /*schedule_modifier=*/nullptr,
-        /*simd_modifier=*/false, /*collapse_val=*/nullptr, /*nowait=*/false,
-        /*ordered_val=*/nullptr, /*order_val=*/nullptr, /*inclusive=*/false);
+        /*simd_modifier=*/false, /*nowait=*/false, /*ordered_val=*/nullptr,
+        /*order_val=*/nullptr, /*inclusive=*/false);
   state.addAttributes(attributes);
 }
 
 LogicalResult WsLoopOp::verify() {
   return verifyReductionVarList(*this, reductions(), reduction_vars());
 }
 
 //===----------------------------------------------------------------------===//

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

convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
 
   // Find the loop configuration.
   llvm::Value *step = moduleTranslation.lookupValue(loop.step()[0]);
   llvm::Type *ivType = step->getType();
   llvm::Value *chunk = nullptr;
   if (loop.schedule_chunk_var()) {
     llvm::Value *chunkVar =
         moduleTranslation.lookupValue(loop.schedule_chunk_var());
-    llvm::Type *chunkVarType = chunkVar->getType();
-    assert(chunkVarType->isIntegerTy() &&
-           "chunk size must be one integer expression");
-    if (chunkVarType->getIntegerBitWidth() < ivType->getIntegerBitWidth())
-      chunk = builder.CreateSExt(chunkVar, ivType);
-    else if (chunkVarType->getIntegerBitWidth() > ivType->getIntegerBitWidth())
-      chunk = builder.CreateTrunc(chunkVar, ivType);
-    else
-      chunk = chunkVar;
+    chunk = builder.CreateSExtOrTrunc(chunkVar, ivType);
   }
 
   SmallVector<omp::ReductionDeclareOp> reductionDecls;
   collectReductionDecls(loop, reductionDecls);
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
       findAllocaInsertPoint(builder, moduleTranslation);
 
   // Allocate space for privatized reduction variables.

mlir/test/Dialect/OpenMP/ops.mlir

  }
 
   return
 }
 
 // CHECK-LABEL: omp_wsloop
 func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
 
-  // CHECK: omp.wsloop collapse(2) ordered(1)
+  // CHECK: omp.wsloop ordered(1)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,0]> : vector<7xi32>, collapse_val = 2, ordered_val = 1} :
+  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,0]> : vector<7xi32>, ordered_val = 1} :
     (index, index, index) -> ()
 
   // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   // CHECK-SAMe: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({
     ^bb0(%iv: index):
       omp.yield
   }) {operand_segment_sizes = dense<[1,1,1,1,1,0,0]> : vector<7xi32>, schedule_val = #omp<"schedulekind static">} :
     (index, index, index, memref<i32>, i32) -> ()
 
   // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>, %{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %linear_var, %linear_var) ({
     ^bb0(%iv: index):
       omp.yield
   }) {operand_segment_sizes = dense<[1,1,1,2,2,0,0]> : vector<7xi32>, schedule_val = #omp<"schedulekind static">} :
     (index, index, index, memref<i32>, memref<i32>, i32, i32) -> ()
 
-  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}) collapse(3) ordered(2)
+  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}) ordered(2)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var, %chunk_var) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,1,1,0,1]> : vector<7xi32>, schedule_val = #omp<"schedulekind dynamic">, collapse_val = 3, ordered_val = 2} :
+  }) {operand_segment_sizes = dense<[1,1,1,1,1,0,1]> : vector<7xi32>, schedule_val = #omp<"schedulekind dynamic">, ordered_val = 2} :
     (index, index, index, memref<i32>, i32, i32) -> ()
 
   // CHECK: omp.wsloop schedule(auto) nowait
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   "omp.wsloop" (%lb, %ub, %step) ({
     ^bb0(%iv: index):
       omp.yield
   }) {operand_segment_sizes = dense<[1,1,1,0,0,0,0]> : vector<7xi32>, nowait, schedule_val = #omp<"schedulekind auto">} :
     (index, index, index) -> ()
 
   return
 }
 
 // CHECK-LABEL: omp_wsloop_pretty
 func.func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32, %chunk_var2 : i16) -> () {
 
-  // CHECK: omp.wsloop collapse(2) ordered(2)
+  // CHECK: omp.wsloop ordered(2)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
-  omp.wsloop collapse(2) ordered(2)
+  omp.wsloop ordered(2)
   for (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
 
   // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   omp.wsloop schedule(static) linear(%data_var = %linear_var : memref<i32>)
   for (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
 
-  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}} : i32) collapse(3) ordered(2)
+  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static = %{{.*}} : i32) ordered(2)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
-  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(static = %chunk_var : i32) collapse(3)
+  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(static = %chunk_var : i32)
   for (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
 
-  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i32, nonmonotonic) collapse(3) ordered(2)
+  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i32, nonmonotonic) ordered(2)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
-  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(dynamic = %chunk_var : i32, nonmonotonic) collapse(3)
+  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(dynamic = %chunk_var : i32, nonmonotonic)
   for (%iv) : index = (%lb) to (%ub) step (%step)  {
     omp.yield
   }
 
-  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i16, monotonic) collapse(3) ordered(2)
+  // CHECK: omp.wsloop linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}} : i16, monotonic) ordered(2)
   // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
-  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(dynamic = %chunk_var2 : i16, monotonic) collapse(3)
+  omp.wsloop ordered(2) linear(%data_var = %linear_var : memref<i32>) schedule(dynamic = %chunk_var2 : i16, monotonic)
   for (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield
   }
 
   // CHECK: omp.wsloop for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
   omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
     omp.yield
   }

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

omp.parallel {
     // CHECK: %[[TOTAL:.*]] = mul nuw i32 %[[PROD]], %[[TRIPCOUNT2]]
     // CHECK: br label %[[COLLAPSED_PREHEADER:.*]]
     //
     // CHECK: [[COLLAPSED_PREHEADER]]:
     // CHECK: store i32 0, ptr
     // CHECK: %[[TOTAL_SUB_1:.*]] = sub i32 %[[TOTAL]], 1
     // CHECK: store i32 %[[TOTAL_SUB_1]], ptr
     // CHECK: call void @__kmpc_for_static_init_4u
-    omp.wsloop collapse(3)
+    omp.wsloop
     for (%arg0, %arg1, %arg2) : i32 = (%0, %1, %2) to (%3, %4, %5) step (%6, %7, %8) {
       %31 = llvm.load %20 : !llvm.ptr<i32>
       %32 = llvm.add %31, %arg0 : i32
       %33 = llvm.add %32, %arg1 : i32
       %34 = llvm.add %33, %arg2 : i32
       llvm.store %34, %20 : !llvm.ptr<i32>
       omp.yield
     }
omp.parallel {
     // CHECK: %[[PROD:.*]] = mul nuw i32 %[[TRIPCOUNT0]], %[[TRIPCOUNT1]]
     // CHECK: %[[TOTAL:.*]] = mul nuw i32 %[[PROD]], %[[TRIPCOUNT2]]
     // CHECK: br label %[[COLLAPSED_PREHEADER:.*]]
     //
     // CHECK: [[COLLAPSED_PREHEADER]]:
     // CHECK: store i32 1, ptr
     // CHECK: store i32 %[[TOTAL]], ptr
     // CHECK: call void @__kmpc_dispatch_init_4u
-    omp.wsloop collapse(3) schedule(dynamic)
+    omp.wsloop schedule(dynamic)
     for (%arg0, %arg1, %arg2) : i32 = (%0, %1, %2) to (%3, %4, %5) step (%6, %7, %8) {
       %31 = llvm.load %20 : !llvm.ptr<i32>
       %32 = llvm.add %31, %arg0 : i32
       %33 = llvm.add %32, %arg1 : i32
       %34 = llvm.add %33, %arg2 : i32
       llvm.store %34, %20 : !llvm.ptr<i32>
       omp.yield
     }