Merge pull request #1798 from schilkp:schilkp/mlir_next_value

PiperOrigin-RevId: 707168188
google · Dec 17, 2024 · 8445840 · 8445840
2 parents ed9dab9 + 12eac1e
commit 8445840
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/xls/contrib/mlir/testdata/ops_translate.mlir b/xls/contrib/mlir/testdata/ops_translate.mlir
@@ -1,4 +1,4 @@
-// RUN: xls_translate --mlir-xls-to-xls %s --main-function=identity
+// RUN: xls_translate --mlir-xls-to-xls %s --main-function=identity | FileCheck --check-prefix=XLS %s
 // This just currently verifies a successful translate. The main function
 // doesn't matter as all are exported, one just set as top.
 
@@ -282,16 +282,19 @@ func.func @constant_scalar() -> i7 {
 
 xls.chan @mychan : i32
 
-// XLS-LABEL: proc eproc({{.*}: bits[32], {{.*}}: (bits[32], bits[1]), {{.*}}: bits[32]}, {{.*}}: bits[16])
-// XLS:  next
+// XLS-LABEL: proc eproc({{.*}}: bits[32], {{.*}}: (bits[32], bits[1]), {{.*}}: bits[1], {{.*}}: (bits[1], bits[8], bits[7]), init={
 xls.eproc @eproc(%arg0: i32 loc("a"), %arg1: tuple<i32, i1> loc("b"),
     %arg2: i1 loc("pred"), %arg3: bf16 loc("fp")) zeroinitializer {
+  // XLS: [[literal:[^ ]*]]: bits[32] = literal(value=6
   %0 = "xls.constant_scalar"() { value = 6 : i32 } : () -> i32
   %tkn1 = "xls.after_all"() : () -> !xls.token
   %tkn_out, %result = xls.blocking_receive %tkn1, @mychan : i32
   %tkn2 = xls.send %tkn_out, %0, @mychan : i32
   %tkn_out2, %result2, %done = xls.nonblocking_receive %tkn2, %arg2, @mychan : i32
-  // CHECK: next_value(param=a, value=a, predicate=pred
-  %2 = xls.next_value [%arg2, %arg0] : (i32) -> i32
+  // XLS: [[not_pred:[^ ]*]]: bits[1] = not(
+  %not_pred = xls.not %arg2 : i1
+  // XLS: next_value(param=a, value=a, predicate=pred
+  // XLS: next_value(param=a, value=[[literal]], predicate=[[not_pred]]
+  %2 = xls.next_value [%arg2, %arg0], [%not_pred, %0] : (i32, i32) -> i32
   xls.yield %2, %arg1, %arg2, %arg3 : i32, tuple<i32, i1>, i1, bf16
 }
diff --git a/xls/contrib/mlir/tools/xls_translate/xls_translate.cc b/xls/contrib/mlir/tools/xls_translate/xls_translate.cc
@@ -1287,7 +1287,7 @@ FailureOr<std::unique_ptr<Package>> mlirXlsToXls(
           auto next_value = cast<NextValueOp>(def);
           for (auto [pred, value] :
                llvm::zip(next_value.getPredicates(), next_value.getValues())) {
-            fb.Next(valueMap[arg], valueMap[yield], /*pred=*/valueMap[pred]);
+            fb.Next(valueMap[arg], valueMap[value], /*pred=*/valueMap[pred]);
           }
         } else {
           fb.Next(valueMap[arg], valueMap[yield]);