Merge pull request #1854 from schilkp:schilkp/mlir_next_value_typing

PiperOrigin-RevId: 716087352
google · Jan 16, 2025 · aaaa243 · aaaa243
2 parents f48427e + 62d6c27
commit aaaa243
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/xls/contrib/mlir/IR/xls_ops.td b/xls/contrib/mlir/IR/xls_ops.td
@@ -1446,10 +1446,10 @@ def Xls_NextValueOp : Xls_Op<"next_value", [Pure, SameVariadicOperandSize]> {
   }];
   let arguments = (ins
     Variadic<I1>:$predicates,
-    Variadic<Xls_Bits>:$values
+    Variadic<Xls_BitsOrTuple>:$values
   );
   let results = (outs
-    Xls_Bits:$result
+    Xls_BitsOrTuple:$result
   );
   // TODO(jpienaar): The verifier should be made stricter (check element type
   // matches).

diff --git a/xls/contrib/mlir/testdata/ops_translate.mlir b/xls/contrib/mlir/testdata/ops_translate.mlir
@@ -298,3 +298,22 @@ xls.eproc @eproc(%arg0: i32 loc("a"), %arg1: tuple<i32, i1> loc("b"),
   %2 = xls.next_value [%arg2, %arg0], [%not_pred, %0] : (i32, i32) -> i32
   xls.yield %2, %arg1, %arg2, %arg3 : i32, tuple<i32, i1>, i1, bf16
 }
+
+// XLS-LABEL: proc eproc2({{.*}}: (bits[32], bits[1]), {{.*}}: bits[1], init={
+xls.eproc @eproc2(%state: tuple<i32, i1> loc("state"), %pred: i1 loc("pred")) zeroinitializer {
+  // XLS: [[literal1:[^ ]*]]: bits[32] = literal(value=6
+  %lit1 = "xls.constant_scalar"() { value = 6 : i32 } : () -> i32
+  // XLS: [[literal2:[^ ]*]]: bits[1] = literal(value=0
+  %lit2 = "xls.constant_scalar"() { value = 0 : i1 } : () -> i1
+  // XLS: [[new_tuple:[^ ]*]]: (bits[32], bits[1]) = tuple(
+  %new_tuple = "xls.tuple"(%lit1, %lit2) : (i32, i1) -> tuple<i32, i1>
+
+  // XLS: [[not_pred:[^ ]*]]: bits[1] = not(
+  %not_pred = xls.not %pred : i1
+
+  // XLS: next_value(param=state, value=state, predicate=pred
+  // XLS: next_value(param=state, value=[[new_tuple]], predicate=[[not_pred]]
+  %next_state = xls.next_value [%pred, %state], [%not_pred, %new_tuple] : (tuple<i32, i1>, tuple<i32, i1>) -> tuple<i32, i1>
+  xls.yield %next_state, %pred : tuple<i32, i1>, i1
+}
+