[RTG] Add interleave_sequences op

Author: maerhart

include/circt/Dialect/RTG/IR/RTGOps.td

@@ -141,6 +141,27 @@ def EmbedSequenceOp : RTGOp<"embed_sequence", []> {
   let assemblyFormat = "$sequence attr-dict";
 }
 
+def InterleaveSequencesOp : RTGOp<"interleave_sequences", [Pure]> {
+  let summary = "interleave a list of sequences";
+  let description = [{
+    This operation takes a list of (at least one) fully randomized sequences and
+    interleaves them by taking the next `batchSize` number of operations
+    implementing the `InstructionOpInterface` of each sequence round-robin.
+
+    Therefore, if only one sequence is in the list, this operation is a NOP.
+  }];
+
+  let arguments = (ins Variadic<RandomizedSequenceType>:$sequences,
+                       DefaultValuedAttr<I32Attr, "1">:$batchSize);
+
+  let results = (outs RandomizedSequenceType:$interleavedSequence);
+
+  let assemblyFormat = "$sequences (`batch` $batchSize^)? attr-dict";
+
+  let hasVerifier = 1;
+  let hasFolder = 1;
+}
+
 //===- Label Operations ---------------------------------------------------===//
 
 class LabelDeclBase<string mnemonic,

lib/Dialect/RTG/IR/RTGOps.cpp

@@ -193,6 +193,24 @@ void SubstituteSequenceOp::print(OpAsmPrinter &p) {
   p.printOptionalAttrDict((*this)->getAttrs(), {});
 }
 
+//===----------------------------------------------------------------------===//
+// InterleaveSequencesOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult InterleaveSequencesOp::verify() {
+  if (getSequences().empty())
+    return emitOpError("must have at least one sequence in the list");
+
+  return success();
+}
+
+OpFoldResult InterleaveSequencesOp::fold(FoldAdaptor adaptor) {
+  if (getSequences().size() == 1)
+    return getSequences()[0];
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // SetCreateOp
 //===----------------------------------------------------------------------===//

test/Dialect/RTG/IR/basic.mlir

@@ -134,3 +134,11 @@ rtg.sequence @integerHandlingOps(%arg0: index, %arg1: index) {
   // CHECK: rtg.random_number_in_range [%arg0, %arg1)
   rtg.random_number_in_range [%arg0, %arg1)
 }
+
+// CHECK-LABEL: rtg.test @interleaveSequences
+rtg.test @interleaveSequences(%seq0: !rtg.randomized_sequence, %seq1: !rtg.randomized_sequence) {
+  // CHECK: rtg.interleave_sequences %seq0 {rtg.some_attr}
+  rtg.interleave_sequences %seq0 {rtg.some_attr}
+  // CHECK: rtg.interleave_sequences %seq0, %seq1 batch 4 {rtg.some_attr}
+  rtg.interleave_sequences %seq0, %seq1 batch 4 {rtg.some_attr}
+}

test/Dialect/RTG/IR/canonicalization.mlir

@@ -0,0 +1,8 @@
+// RUN: circt-opt --canonicalize %s | FileCheck %s
+
+// CHECK-LABEL: @interleaveSequences
+rtg.test @interleaveSequences(%seq0: !rtg.randomized_sequence) {
+  // CHECK-NEXT: rtg.embed_sequence %seq0
+  %0 = rtg.interleave_sequences %seq0
+  rtg.embed_sequence %0
+}

test/Dialect/RTG/IR/errors.mlir

@@ -174,3 +174,10 @@ rtg.target @target : !rtg.dict<> {
   // expected-error @below {{third sequence element type must be a fully substituted sequence}}
   rtg.context_switch #rtgtest.cpu<0> -> #rtgtest.cpu<1>, %0 : !rtg.sequence<!rtgtest.cpu, !rtgtest.cpu, !rtgtest.cpu>
 }
+
+// -----
+
+rtg.test @test() {
+  // expected-error @below {{must have at least one sequence in the list}}
+  %0 = rtg.interleave_sequences
+}