dslx: Add max_trace_verbosity flag to DSLX interpreter

Author: joannabrozek

xls/dslx/bytecode/bytecode.h

@@ -330,21 +330,26 @@ class Bytecode {
     TraceData& operator=(TraceData&& other) = default;
 
     TraceData(std::vector<FormatStep> steps,
-              std::vector<ValueFormatDescriptor> value_fmt_descs)
+              std::vector<ValueFormatDescriptor> value_fmt_descs,
+              std::optional<int64_t> verbosity = std::nullopt)
         : steps_(std::move(steps)),
-          value_fmt_descs_(std::move(value_fmt_descs)) {}
+          value_fmt_descs_(std::move(value_fmt_descs)),
+          verbosity_(verbosity) {}
 
     absl::Span<const FormatStep> steps() const { return steps_; }
     absl::Span<const ValueFormatDescriptor> value_fmt_descs() const {
       return value_fmt_descs_;
     }
+    const std::optional<int64_t> verbosity() const { return verbosity_; }
 
    private:
     std::vector<FormatStep> steps_;
 
     // For default formatting of struct operands we hold metadata that allows us
     // to format them in more detail (struct name, fields, etc).
     std::vector<ValueFormatDescriptor> value_fmt_descs_;
+    // Verbosity of the `vtrace` messages
+    std::optional<int64_t> verbosity_;
   };
 
   // Information necessary for channel operations.

xls/dslx/bytecode/bytecode_emitter.cc

@@ -1016,9 +1016,19 @@ absl::Status BytecodeEmitter::HandleFormatMacro(const FormatMacro* node) {
 
   CHECK(node->macro() == "trace_fmt!" || node->macro() == "vtrace_fmt!")
       << "Expected trace_fmt! or vtrace_fmt! but got: " << node->macro();
+  std::optional<int64_t> verbosity = std::nullopt;
+  if (node->verbosity().has_value()) {
+    XLS_ASSIGN_OR_RETURN(InterpValue verbosity_val,
+                         type_info_->GetConstExpr(*node->verbosity()));
+    absl::StatusOr<int64_t> verbosity_as_int64 =
+        verbosity_val.GetBitValueViaSign();
+    if (verbosity_as_int64.ok()) {
+      verbosity = *verbosity_as_int64;
+    }
+  }
   Bytecode::TraceData trace_data(
       std::vector<FormatStep>(node->format().begin(), node->format().end()),
-      std::move(value_fmt_descs));
+      std::move(value_fmt_descs), verbosity);
   bytecode_.push_back(
       Bytecode(node->span(), Bytecode::Op::kTraceFmt, std::move(trace_data)));
   return absl::OkStatus();

xls/dslx/bytecode/bytecode_interpreter.cc

@@ -1437,9 +1437,13 @@ absl::Status BytecodeInterpreter::EvalTraceFmt(const Bytecode& bytecode) {
                        bytecode.trace_data());
   XLS_ASSIGN_OR_RETURN(std::string message,
                        TraceDataToString(*trace_data, stack_));
-  if (options_.trace_hook()) {
+  // Exclude "vtrace" messages if the verbosity level is set to a value higher
+  // than the verbosity level passed by the "max_trace_verbosity" flag.
+  if (options_.trace_hook() &&
+      (*trace_data->verbosity() <= *options_.max_trace_verbosity())) {
     options_.trace_hook()(bytecode.source_span(), message);
   }
+
   stack_.Push(InterpValue::MakeToken());
   return absl::OkStatus();
 }

xls/dslx/bytecode/bytecode_interpreter_options.h

@@ -100,6 +100,18 @@ class BytecodeInterpreterOptions {
   }
   std::optional<int64_t> max_ticks() const { return max_ticks_; }
 
+  // Store verbosity level from the `max_trace_verbosity` flag. Used to filter
+  // `vtrace` messages if their verbosity level is set higher than the verbosity
+  // level passed by the flag.
+  BytecodeInterpreterOptions& max_trace_verbosity(
+      std::optional<int64_t> value) {
+    max_trace_verbosity_ = value;
+    return *this;
+  }
+  std::optional<int64_t> max_trace_verbosity() const {
+    return max_trace_verbosity_;
+  }
+
   void set_validate_final_stack_depth(bool enabled) {
     validate_final_stack_depth_ = enabled;
   }
@@ -123,6 +135,7 @@ class BytecodeInterpreterOptions {
   bool trace_channels_ = false;
   bool trace_calls_ = false;
   std::optional<int64_t> max_ticks_;
+  std::optional<int64_t> max_trace_verbosity_;
   bool validate_final_stack_depth_ = true;
   FormatPreference format_preference_ = FormatPreference::kDefault;
 };

xls/dslx/bytecode/bytecode_interpreter_test.cc

@@ -226,6 +226,28 @@ fn main() -> () {
   EXPECT_EQ(value, InterpValue::MakeUnit());
 }
 
+TEST_F(BytecodeInterpreterTest, VtraceFmtBitsValueDefaultFormat) {
+  constexpr std::string_view kProgram = R"(
+fn main() -> () {
+  vtrace_fmt!(u32:2, "{}", u32:2);
+  vtrace_fmt!(u32:6, "{}", u32:6);
+  vtrace_fmt!(u32:10, "{}", u32:10);
+  vtrace_fmt!(u32:16, "{}", u32:16);
+}
+)";
+  std::vector<std::string> trace_output;
+  XLS_ASSERT_OK_AND_ASSIGN(
+      InterpValue value,
+      Interpret(kProgram, "main", /*args=*/{},
+                BytecodeInterpreterOptions()
+                    .trace_hook([&](const Span&, std::string_view s) {
+                      trace_output.push_back(std::string{s});
+                    })
+                    .max_trace_verbosity(8)));
+  EXPECT_THAT(trace_output, ElementsAre("2", "6"));
+  EXPECT_EQ(value, InterpValue::MakeUnit());
+}
+
 TEST_F(BytecodeInterpreterTest, TraceFmtBitsValueHexFormat) {
   constexpr std::string_view kProgram = R"(
 fn main() -> () {

xls/dslx/interpreter_main.cc

@@ -89,6 +89,9 @@ ABSL_FLAG(bool, trace_calls, false,
 ABSL_FLAG(int64_t, max_ticks, 100000,
           "If non-zero, the maximum number of ticks to execute on any proc. If "
           "exceeded an error is returned.");
+ABSL_FLAG(int64_t, max_trace_verbosity, 0,
+          "Maximum verbosity for traces. Traces with higher verbosity are "
+          "stripped from the output. 0 by default.");
 ABSL_FLAG(std::string, evaluator, "dslx-interpreter",
           "What evaluator should be used to actually execute the dslx test. "
           "'dslx-interpreter' is the DSLX bytecode interpreter. 'ir-jit' is "
@@ -169,6 +172,11 @@ absl::StatusOr<TestResult> RealMain(
                                    : TypeInferenceVersion::kVersion1)
           : std::nullopt;
 
+  std::optional<int64_t> max_trace_verbosity =
+      absl::GetFlag(FLAGS_max_trace_verbosity) == 0
+          ? std::nullopt
+          : std::optional<int64_t>(absl::GetFlag(FLAGS_max_trace_verbosity));
+
   RealFilesystem vfs;
 
   XLS_ASSIGN_OR_RETURN(std::string program,
@@ -243,6 +251,7 @@ absl::StatusOr<TestResult> RealMain(
       .trace_channels = trace_channels,
       .trace_calls = trace_calls,
       .max_ticks = max_ticks,
+      .max_trace_verbosity = max_trace_verbosity,
   };
 
   std::unique_ptr<AbstractTestRunner> test_runner = GetTestRunner(evaluator);

xls/dslx/run_routines/run_routines.cc

@@ -1025,6 +1025,7 @@ absl::StatusOr<TestResultData> AbstractTestRunner::ParseAndTest(
         .trace_channels(options.trace_channels)
         .trace_calls(options.trace_calls)
         .max_ticks(options.max_ticks)
+        .max_trace_verbosity(options.max_trace_verbosity)
         .format_preference(options.format_preference);
     if (std::holds_alternative<TestFunction*>(*member)) {
       XLS_ASSIGN_OR_RETURN(

xls/dslx/run_routines/run_routines.h

@@ -109,6 +109,7 @@ struct ParseAndTestOptions {
   bool trace_channels = false;
   bool trace_calls = false;
   std::optional<int64_t> max_ticks;
+  std::optional<int64_t> max_trace_verbosity;
   std::function<std::unique_ptr<VirtualizableFilesystem>()> vfs_factory =
       nullptr;
 };