[MLIR][OpenMP] Add omp.simd operationupstream/users/skatrak/spr/simd-mlir

This patch introduces the `omp.simd` operation. In contrast to the existing
`omp.simdloop` operation, it is intended to hold SIMD information within
worksharing loops, rather than representing a SIMD-only loop. Some examples
of such loops are "omp do/for simd", "omp distribute simd", "omp target teams
distribute parallel do/for simd", etc. For more context on this work, refer to
PR #79559.

This operation must always be nested within an `omp.wsloop` operation as its
only non-terminator child. It follows the same approach as the `omp.distribute`
operation, by serving as a simple wrapper operation holding clause information.

4 files changed, 515 insertions, 12 deletions

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 96c15e775a30..50b316a21554 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -505,6 +505,9 @@ def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
 
     /// Returns the number of reduction variables.
     unsigned getNumReductionVars() { return getReductionVars().size(); }
+
+    /// Returns its nested 'omp.simd' operation, if present.
+    SimdOp getNestedSimd();
   }];
   let hasCustomAssemblyFormat = 1;
   let assemblyFormat = [{
@@ -617,11 +620,84 @@ def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
   let hasVerifier = 1;
 }
 
+def SimdOp : OpenMP_Op<"simd",
+    [AttrSizedOperandSegments, MemoryEffects<[MemWrite]>,
+     HasParent<"WsLoopOp">]> {
+ let summary = "simd construct";
+  let description = [{
+    The simd construct can be applied to a loop to indicate that the loop can be
+    transformed into a SIMD loop (that is, multiple iterations of the loop can
+    be executed concurrently using SIMD instructions).
+    
+    This operation is intended to hold SIMD information for a worksharing loop
+    (i.e. "omp for simd"), so it must always be nested inside of a parent
+    "omp.wsloop" operation as its only child. For SIMD loops not combined with a
+    worksharing loop (i.e. "omp simd"), the "omp.simdloop" is used instead.
+
+    The body region can contain any number of blocks. The region is terminated
+    by "omp.yield" instruction without operands.
+
+    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
+    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.wsloop for (%i) : index = (%c0) to (%c10) step (%c1) {
+      omp.simd <clauses> {
+        // block operations
+        omp.yield
+      }
+      omp.yield
+    ```
+  }];
+
+  // TODO: Add other clauses
+  let arguments = (ins 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
+     );
+
+  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  `)`
+    ) $region attr-dict
+  }];
+
+  let hasVerifier = 1;
+}
 
 def YieldOp : OpenMP_Op<"yield",
     [Pure, ReturnLike, Terminator,
      ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
-     "AtomicUpdateOp", "SimdLoopOp"]>]> {
+     "AtomicUpdateOp", "SimdLoopOp", "SimdOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 13cc16125a27..caa888d030f7 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1131,28 +1131,33 @@ void printLoopControl(OpAsmPrinter &p, Operation *op, Region &region,
 }
 
 //===----------------------------------------------------------------------===//
-// Verifier for Simd construct [2.9.3.1]
+// Verifier for Simd constructs [2.9.3.1]
 //===----------------------------------------------------------------------===//
 
-LogicalResult SimdLoopOp::verify() {
-  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()
+template <typename OpTy>
+static LogicalResult verifySimdOp(OpTy op) {
+  if (op.getSimdlen().has_value() && op.getSafelen().has_value() &&
+      op.getSimdlen().value() > op.getSafelen().value()) {
+    return op.emitOpError()
            << "simdlen clause and safelen clause are both present, but the "
               "simdlen value is not less than or equal to safelen value";
   }
-  if (verifyAlignedClause(*this, this->getAlignmentValues(),
-                          this->getAlignedVars())
+  if (verifyAlignedClause(op, op.getAlignmentValues(), op.getAlignedVars())
           .failed())
     return failure();
-  if (verifyNontemporalClause(*this, this->getNontemporalVars()).failed())
+  if (verifyNontemporalClause(op, op.getNontemporalVars()).failed())
     return failure();
   return success();
 }
 
+LogicalResult SimdLoopOp::verify() {
+  if (this->getLowerBound().empty())
+    return emitOpError() << "empty lowerbound for simd loop operation";
+  return verifySimdOp(*this);
+}
+
+LogicalResult SimdOp::verify() { return verifySimdOp(*this); }
+
 //===----------------------------------------------------------------------===//
 // Verifier for Distribute construct [2.9.4.1]
 //===----------------------------------------------------------------------===//
@@ -1329,7 +1334,34 @@ void WsLoopOp::build(OpBuilder &builder, OperationState &state,
   state.addAttributes(attributes);
 }
 
+SimdOp WsLoopOp::getNestedSimd() {
+  auto ops = this->getOps<SimdOp>();
+  assert(std::distance(ops.begin(), ops.end()) <= 1 &&
+         "There can only be a single omp.simd child at most");
+  return ops.empty() ? SimdOp() : *ops.begin();
+}
+
 LogicalResult WsLoopOp::verify() {
+  // Check that, if it has an omp.simd child, it must be the only one.
+  bool hasSimd = false, hasOther = false;
+  for (auto &op : this->getOps()) {
+    if (isa<SimdOp>(op)) {
+      if (hasSimd)
+        return emitOpError() << "cannot have multiple 'omp.simd' child ops";
+      hasSimd = true;
+
+      if (hasOther)
+        break;
+    } else if (!op.hasTrait<OpTrait::IsTerminator>()) {
+      hasOther = true;
+      if (hasSimd)
+        break;
+    }
+  }
+  if (hasSimd && hasOther)
+    return emitOpError() << "if 'omp.simd' is a child, it must be the only "
+                            "non-terminator child op";
+
   return verifyReductionVarList(*this, getReductions(), getReductionVars());
 }
 
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 812b79e35595..29a6a078fad8 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -192,6 +192,50 @@ llvm.func @test_omp_wsloop_dynamic_wrong_modifier3(%lb : i64, %ub : i64, %step :
 
 // -----
 
+llvm.func @test_omp_wsloop_simd_multiple(%lb : i64, %ub : i64, %step : i64) -> () {
+  // expected-error @+1 {{op cannot have multiple 'omp.simd' child ops}}
+  omp.wsloop for (%iv) : i64 = (%lb) to (%ub) step (%step) {
+    omp.simd {
+      omp.yield
+    }
+    omp.simd {
+      omp.yield
+    }
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @test_omp_wsloop_simd_invalid_before(%lb : i64, %ub : i64, %step : i64) -> () {
+  // expected-error @+1 {{op if 'omp.simd' is a child, it must be the only non-terminator child op}}
+  omp.wsloop for (%iv) : i64 = (%lb) to (%ub) step (%step) {
+    %c1 = arith.constant 1 : i32
+    omp.simd {
+      omp.yield
+    }
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
+llvm.func @test_omp_wsloop_simd_invalid_after(%lb : i64, %ub : i64, %step : i64) -> () {
+  // expected-error @+1 {{op if 'omp.simd' is a child, it must be the only non-terminator child op}}
+  omp.wsloop for (%iv) : i64 = (%lb) to (%ub) step (%step) {
+    omp.simd {
+      omp.yield
+    }
+    %c1 = arith.constant 1 : i32
+    omp.yield
+  }
+  llvm.return
+}
+
+// -----
+
 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) ({
@@ -205,6 +249,18 @@ func.func @omp_simdloop(%lb : index, %ub : index, %step : i32) -> () {
 
 // -----
 
+func.func @omp_simd(%lb : index, %ub : index, %step : i32) -> () {
+  // expected-error @below {{'omp.simd' op expects parent op 'omp.wsloop'}}
+  "omp.simd" () ({
+    ^bb0(%iv: index):
+      omp.yield
+  }) {operandSegmentSizes = array<i32: 0, 0, 0>} : () -> ()
+
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index,
                                        %data_var : memref<i32>) -> () {
   //  expected-error @below {{expected '->'}}
@@ -217,6 +273,20 @@ func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index,
 
 // -----
 
+func.func @omp_simd_pretty_aligned(%lb : index, %ub : index, %step : index,
+                                   %data_var : memref<i32>) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    //  expected-error @below {{expected '->'}}
+    omp.simd aligned(%data_var : memref<i32>) {
+      omp.yield
+    }
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_aligned_mismatch(%arg0 : index, %arg1 : index,
                                          %arg2 : index, %arg3 : memref<i32>,
                                          %arg4 : memref<i32>) -> () {
@@ -231,6 +301,22 @@ func.func @omp_simdloop_aligned_mismatch(%arg0 : index, %arg1 : index,
 
 // -----
 
+func.func @omp_simd_aligned_mismatch(%arg0 : index, %arg1 : index,
+                                     %arg2 : index, %arg3 : memref<i32>,
+                                     %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    //  expected-error @below {{op expected as many alignment values as aligned variables}}
+    "omp.simd"(%arg3, %arg4) ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [128],
+        operandSegmentSizes = array<i32: 2, 0, 0>} : (memref<i32>, memref<i32>) -> ()
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_aligned_negative(%arg0 : index, %arg1 : index,
                                          %arg2 : index, %arg3 : memref<i32>,
                                          %arg4 : memref<i32>) -> () {
@@ -244,6 +330,21 @@ func.func @omp_simdloop_aligned_negative(%arg0 : index, %arg1 : index,
 
 // -----
 
+func.func @omp_simd_aligned_negative(%arg0 : index, %arg1 : index,
+                                     %arg2 : index, %arg3 : memref<i32>,
+                                     %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    //  expected-error @below {{op alignment should be greater than 0}}
+    "omp.simd"(%arg3, %arg4) ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [-1, 128], operandSegmentSizes = array<i32: 2, 0, 0>} : (memref<i32>, memref<i32>) -> ()
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_unexpected_alignment(%arg0 : index, %arg1 : index,
                                              %arg2 : index, %arg3 : memref<i32>,
                                              %arg4 : memref<i32>) -> () {
@@ -257,6 +358,21 @@ func.func @omp_simdloop_unexpected_alignment(%arg0 : index, %arg1 : index,
 
 // -----
 
+func.func @omp_simd_unexpected_alignment(%arg0 : index, %arg1 : index,
+                                             %arg2 : index, %arg3 : memref<i32>,
+                                             %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    //  expected-error @below {{unexpected alignment values attribute}}
+    "omp.simd"() ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [1, 128], operandSegmentSizes = array<i32: 0, 0, 0>} : () -> ()
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_aligned_float(%arg0 : index, %arg1 : index,
                                       %arg2 : index, %arg3 : memref<i32>,
                                       %arg4 : memref<i32>) -> () {
@@ -270,6 +386,21 @@ func.func @omp_simdloop_aligned_float(%arg0 : index, %arg1 : index,
 
 // -----
 
+func.func @omp_simd_aligned_float(%arg0 : index, %arg1 : index,
+                                  %arg2 : index, %arg3 : memref<i32>,
+                                  %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    //  expected-error @below {{failed to satisfy constraint: 64-bit integer array attribute}}
+    "omp.simd"(%arg3, %arg4) ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [1.5, 128], operandSegmentSizes = array<i32: 2, 0, 0>} : (memref<i32>, memref<i32>) -> ()
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_aligned_the_same_var(%arg0 : index, %arg1 : index,
                                              %arg2 : index, %arg3 : memref<i32>,
                                              %arg4 : memref<i32>) -> () {
@@ -283,6 +414,21 @@ func.func @omp_simdloop_aligned_the_same_var(%arg0 : index, %arg1 : index,
 
 // -----
 
+func.func @omp_simd_aligned_the_same_var(%arg0 : index, %arg1 : index,
+                                         %arg2 : index, %arg3 : memref<i32>,
+                                         %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    //  expected-error @below {{aligned variable used more than once}}
+    "omp.simd"(%arg3, %arg3) ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [1, 128], operandSegmentSizes = array<i32: 2, 0, 0>} : (memref<i32>, memref<i32>) -> ()
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_nontemporal_the_same_var(%arg0 : index,
                                                  %arg1 : index,
                                                  %arg2 : index,
@@ -297,6 +443,22 @@ func.func @omp_simdloop_nontemporal_the_same_var(%arg0 : index,
 
 // -----
 
+func.func @omp_simd_nontemporal_the_same_var(%arg0 : index,
+                                             %arg1 : index,
+                                             %arg2 : index,
+                                             %arg3 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    //  expected-error @below {{nontemporal variable used more than once}}
+    "omp.simd"(%arg3, %arg3) ({
+      "omp.yield"() : () -> ()
+    }) {operandSegmentSizes = array<i32: 0, 0, 2>} : (memref<i32>, memref<i32>) -> ()
+    omp.yield
+  }
+  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) {
@@ -307,6 +469,19 @@ func.func @omp_simdloop_order_value(%lb : index, %ub : index, %step : index) {
 
 // -----
 
+func.func @omp_simd_order_value(%lb : index, %ub : index, %step : index) {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // expected-error @below {{invalid clause value: 'default'}}
+    omp.simd order(default) {
+      omp.yield
+    }
+    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) {
@@ -317,6 +492,19 @@ func.func @omp_simdloop_pretty_simdlen(%lb : index, %ub : index, %step : index)
 
 // -----
 
+func.func @omp_simd_pretty_simdlen(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // expected-error @below {{op attribute 'simdlen' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive}}
+    omp.simd simdlen(0) {
+      omp.yield
+    }
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_pretty_safelen(%lb : index, %ub : index, %step : index) -> () {
   // expected-error @below {{op attribute 'safelen' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive}}
   omp.simdloop safelen(0) for (%iv): index = (%lb) to (%ub) step (%step) {
@@ -327,6 +515,19 @@ func.func @omp_simdloop_pretty_safelen(%lb : index, %ub : index, %step : index)
 
 // -----
 
+func.func @omp_simd_pretty_safelen(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // expected-error @below {{op attribute 'safelen' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive}}
+    omp.simd safelen(0) {
+      omp.yield
+    }
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 func.func @omp_simdloop_pretty_simdlen_safelen(%lb : index, %ub : index, %step : index) -> () {
   // expected-error @below {{'omp.simdloop' op simdlen clause and safelen clause are both present, but the simdlen value is not less than or equal to safelen value}}
   omp.simdloop simdlen(2) safelen(1) for (%iv): index = (%lb) to (%ub) step (%step) {
@@ -337,6 +538,19 @@ func.func @omp_simdloop_pretty_simdlen_safelen(%lb : index, %ub : index, %step :
 
 // -----
 
+func.func @omp_simd_pretty_simdlen_safelen(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // expected-error @below {{'omp.simd' op simdlen clause and safelen clause are both present, but the simdlen value is not less than or equal to safelen value}}
+    omp.simd simdlen(2) safelen(1) {
+      omp.yield
+    }
+    omp.yield
+  }
+  return
+}
+
+// -----
+
 // expected-error @below {{op expects initializer region with one argument of the reduction type}}
 omp.reduction.declare @add_f32 : f64
 init {
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index ccf72ae31d43..2487f0eb5654 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -176,6 +176,21 @@ func.func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memre
   }) {operandSegmentSizes = array<i32: 1,1,1,0,0,0,0>, nowait, schedule_val = #omp<schedulekind auto>} :
     (index, index, index) -> ()
 
+  // CHECK: omp.wsloop
+  // CHECK-SAME: for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
+  // CHECK-NEXT: omp.simd
+  // CHECK-NEXT: omp.yield
+  // CHECK-NEXT: }
+  // CHECK-NEXT: omp.yield
+  "omp.wsloop" (%lb, %ub, %step) ({
+    ^bb0(%iv: index):
+      "omp.simd" () ({
+        omp.yield
+      }) : () -> ()
+      omp.yield
+  }) {operandSegmentSizes = array<i32: 1,1,1,0,0,0,0>} :
+    (index, index, index) -> ()
+
   return
 }
 
@@ -339,6 +354,19 @@ func.func @omp_simdloop(%lb : index, %ub : index, %step : index) -> () {
   return
 }
 
+// CHECK-LABEL: omp_simd
+func.func @omp_simd(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd {
+    "omp.simd" () ({
+      omp.yield
+    }) {operandSegmentSizes = array<i32: 0,0,0>} : () -> ()
+    omp.yield
+  }
+
+  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>) -> () {
@@ -353,6 +381,21 @@ func.func @omp_simdloop_aligned_list(%arg0 : index, %arg1 : index, %arg2 : index
   return
 }
 
+// CHECK-LABEL: omp_simd_aligned_list
+func.func @omp_simd_aligned_list(%arg0 : index, %arg1 : index, %arg2 : index,
+                                 %arg3 : memref<i32>, %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    // CHECK:      omp.simd   aligned(%{{.*}} : memref<i32> -> 32 : i64,
+    // CHECK-SAME: %{{.*}} : memref<i32> -> 128 : i64) {
+    "omp.simd"(%arg3, %arg4) ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [32, 128],
+        operandSegmentSizes = array<i32: 2, 0, 0>} : (memref<i32>, memref<i32>) -> ()
+    omp.yield
+  }
+  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>) -> () {
@@ -366,6 +409,20 @@ func.func @omp_simdloop_aligned_single(%arg0 : index, %arg1 : index, %arg2 : ind
   return
 }
 
+// CHECK-LABEL: omp_simd_aligned_single
+func.func @omp_simd_aligned_single(%arg0 : index, %arg1 : index, %arg2 : index,
+                                   %arg3 : memref<i32>, %arg4 : memref<i32>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    // CHECK:      omp.simd   aligned(%{{.*}} : memref<i32> -> 32 : i64) {
+    "omp.simd"(%arg3) ({
+      "omp.yield"() : () -> ()
+    }) {alignment_values = [32],
+        operandSegmentSizes = array<i32: 1, 0, 0>} : (memref<i32>) -> ()
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_nontemporal_list
 func.func @omp_simdloop_nontemporal_list(%arg0 : index,
                                          %arg1 : index,
@@ -381,6 +438,22 @@ func.func @omp_simdloop_nontemporal_list(%arg0 : index,
   return
 }
 
+// CHECK-LABEL: omp_simd_nontemporal_list
+func.func @omp_simd_nontemporal_list(%arg0 : index,
+                                     %arg1 : index,
+                                     %arg2 : index,
+                                     %arg3 : memref<i32>,
+                                     %arg4 : memref<i64>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    // CHECK:      omp.simd   nontemporal(%{{.*}}, %{{.*}} : memref<i32>, memref<i64>) {
+    "omp.simd"(%arg3, %arg4) ({
+      "omp.yield"() : () -> ()
+    }) {operandSegmentSizes = array<i32: 0, 0, 2>} : (memref<i32>, memref<i64>) -> ()
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_nontemporal_single
 func.func @omp_simdloop_nontemporal_single(%arg0 : index,
                                            %arg1 : index,
@@ -396,6 +469,22 @@ func.func @omp_simdloop_nontemporal_single(%arg0 : index,
   return
 }
 
+// CHECK-LABEL: omp_simd_nontemporal_single
+func.func @omp_simd_nontemporal_single(%arg0 : index,
+                                           %arg1 : index,
+                                           %arg2 : index,
+                                           %arg3 : memref<i32>,
+                                           %arg4 : memref<i64>) -> () {
+  omp.wsloop for (%arg5) : index = (%arg0) to (%arg1) step (%arg2) {
+    // CHECK:      omp.simd   nontemporal(%{{.*}} : memref<i32>) {
+    "omp.simd"(%arg3) ({
+      "omp.yield"() : () -> ()
+    }) {operandSegmentSizes = array<i32: 0, 0, 1>} : (memref<i32>) -> ()
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_pretty
 func.func @omp_simdloop_pretty(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
@@ -405,6 +494,18 @@ func.func @omp_simdloop_pretty(%lb : index, %ub : index, %step : index) -> () {
   return
 }
 
+// CHECK-LABEL: omp_simd_pretty
+func.func @omp_simd_pretty(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd
+    omp.simd {
+      omp.yield
+    }
+    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>,
@@ -419,6 +520,21 @@ func.func @omp_simdloop_pretty_aligned(%lb : index, %ub : index, %step : index,
   return
 }
 
+// CHECK-LABEL:   func.func @omp_simd_pretty_aligned(
+func.func @omp_simd_pretty_aligned(%lb : index, %ub : index, %step : index,
+                                   %data_var : memref<i32>,
+                                   %data_var1 : memref<i32>) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK:      omp.simd   aligned(%{{.*}} : memref<i32> -> 32 : i64,
+    // CHECK-SAME: %{{.*}} : memref<i32> -> 128 : i64) {
+    omp.simd aligned(%data_var :  memref<i32> -> 32, %data_var1 : memref<i32> -> 128) {
+      omp.yield
+    }
+    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 (%{{.*}})
@@ -428,6 +544,18 @@ func.func @omp_simdloop_pretty_if(%lb : index, %ub : index, %step : index, %if_c
   return
 }
 
+// CHECK-LABEL: omp_simd_pretty_if
+func.func @omp_simd_pretty_if(%lb : index, %ub : index, %step : index, %if_cond : i1) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd if(%{{.*}}) {
+    omp.simd if(%if_cond) {
+      omp.yield
+    }
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL:   func.func @omp_simdloop_pretty_nontemporal
 func.func @omp_simdloop_pretty_nontemporal(%lb : index,
                                            %ub : index,
@@ -442,6 +570,23 @@ func.func @omp_simdloop_pretty_nontemporal(%lb : index,
   }
   return
 }
+
+// CHECK-LABEL:   func.func @omp_simd_pretty_nontemporal
+func.func @omp_simd_pretty_nontemporal(%lb : index,
+                                           %ub : index,
+                                           %step : index,
+                                           %data_var : memref<i32>,
+                                           %data_var1 : memref<i32>) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd   nontemporal(%{{.*}}, %{{.*}} : memref<i32>, memref<i32>) {
+    omp.simd nontemporal(%data_var, %data_var1 : memref<i32>, memref<i32>) {
+      omp.yield
+    }
+    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)
@@ -452,6 +597,18 @@ func.func @omp_simdloop_pretty_order(%lb : index, %ub : index, %step : index) ->
   return
 }
 
+// CHECK-LABEL: omp_simd_pretty_order
+func.func @omp_simd_pretty_order(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd order(concurrent) {
+    omp.simd order(concurrent) {
+      omp.yield
+    }
+    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 (%{{.*}})
@@ -461,6 +618,18 @@ func.func @omp_simdloop_pretty_simdlen(%lb : index, %ub : index, %step : index)
   return
 }
 
+// CHECK-LABEL: omp_simd_pretty_simdlen
+func.func @omp_simd_pretty_simdlen(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd simdlen(2) {
+    omp.simd simdlen(2) {
+      omp.yield
+    }
+    omp.yield
+  }
+  return
+}
+
 // CHECK-LABEL: omp_simdloop_pretty_safelen
 func.func @omp_simdloop_pretty_safelen(%lb : index, %ub : index, %step : index) -> () {
   // CHECK: omp.simdloop safelen(2) for (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}})
@@ -470,6 +639,18 @@ func.func @omp_simdloop_pretty_safelen(%lb : index, %ub : index, %step : index)
   return
 }
 
+// CHECK-LABEL: omp_simd_pretty_safelen
+func.func @omp_simd_pretty_safelen(%lb : index, %ub : index, %step : index) -> () {
+  omp.wsloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+    // CHECK: omp.simd safelen(2) {
+    omp.simd safelen(2) {
+      omp.yield
+    }
+    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 (%{{.*}}, %{{.*}})