AOT block wrapper

This adds BLOCK type support to `cc_xls_ir_jit_wrapper` targets. These wrappers include helper methods to read each output port (including as C++ values if they are a standard c++ integer size) and a <Wrapper>Ports struct which can be used to set each port (again as either value or c++). Fixes: #1651 PiperOrigin-RevId: 687421784
google · Oct 18, 2024 · b2fe633 · b2fe633
1 parent 84705c8
commit b2fe633
Show file tree

Hide file tree

Showing 15 changed files with 800 additions and 83 deletions.
diff --git a/docs_src/bazel_rules_macros.md b/docs_src/bazel_rules_macros.md
@@ -608,7 +608,7 @@ currently produced should be considered INCOMPLETE.
 ## cc_xls_ir_jit_wrapper
 
 <pre>
-cc_xls_ir_jit_wrapper(<a href="#cc_xls_ir_jit_wrapper-name">name</a>, <a href="#cc_xls_ir_jit_wrapper-src">src</a>, <a href="#cc_xls_ir_jit_wrapper-jit_wrapper_args">jit_wrapper_args</a>, <a href="#cc_xls_ir_jit_wrapper-wrapper_type">wrapper_type</a>, <a href="#cc_xls_ir_jit_wrapper-kwargs">kwargs</a>)
+cc_xls_ir_jit_wrapper(<a href="#cc_xls_ir_jit_wrapper-name">name</a>, <a href="#cc_xls_ir_jit_wrapper-src">src</a>, <a href="#cc_xls_ir_jit_wrapper-jit_wrapper_args">jit_wrapper_args</a>, <a href="#cc_xls_ir_jit_wrapper-wrapper_type">wrapper_type</a>, <a href="#cc_xls_ir_jit_wrapper-top">top</a>, <a href="#cc_xls_ir_jit_wrapper-kwargs">kwargs</a>)
 </pre>
 
 Invokes the JIT wrapper generator and compiles the result as a cc_library.
@@ -627,6 +627,7 @@ identical to this macro.
 | <a id="cc_xls_ir_jit_wrapper-src"></a>src |  The path to the IR file.   |  none |
 | <a id="cc_xls_ir_jit_wrapper-jit_wrapper_args"></a>jit_wrapper_args |  Arguments of the JIT wrapper tool. Note: argument 'output_name' cannot be defined.   |  `{}` |
 | <a id="cc_xls_ir_jit_wrapper-wrapper_type"></a>wrapper_type |  The type of XLS construct to wrap. Must be one of 'BLOCK', 'FUNCTION', or 'PROC'. You should use the exported FUNCTION_WRAPPER_TYPE, BLOCK_WRAPPER_TYPE, or PROC_WRAPPER_TYPE symbols. Defaults to FUNCTION_WRAPPER_TYPE for compatibility.   |  `"FUNCTION"` |
+| <a id="cc_xls_ir_jit_wrapper-top"></a>top |  Name of the top function/proc/block.   |  `""` |
 | <a id="cc_xls_ir_jit_wrapper-kwargs"></a>kwargs |  Keyword arguments. Named arguments.   |  none |
 
 

diff --git a/xls/build_rules/tests/BUILD b/xls/build_rules/tests/BUILD
@@ -607,6 +607,7 @@ xls_ir_jit_wrapper_macro(
         "namespace": "not_xls::test",
     },
     source_file = "simple_example_5_one_stage_jit_wrapper.cc",
+    top = "main",
     wrapper_type = FUNCTION_WRAPPER_TYPE,
 )
 
@@ -826,6 +827,7 @@ xls_ir_jit_wrapper_macro(
         "namespace": "not_xls::test",
     },
     source_file = "xls_ir_jit_wrapper.cc",
+    top = "main",
     wrapper_type = FUNCTION_WRAPPER_TYPE,
 )
 

diff --git a/xls/build_rules/xls_jit_wrapper_rules.bzl b/xls/build_rules/xls_jit_wrapper_rules.bzl
@@ -36,7 +36,11 @@ load(
     "//xls/build_rules:xls_internal_build_defs.bzl",
     "XLS_IS_MSAN_BUILD",
 )
-load("//xls/build_rules:xls_ir_rules.bzl", "xls_ir_common_attrs")
+load(
+    "//xls/build_rules:xls_ir_rules.bzl",
+    "xls_ir_common_attrs",
+    "xls_ir_top_attrs",
+)
 load("//xls/build_rules:xls_providers.bzl", "AotCompileInfo", "JitWrapperInfo")
 load(
     "//xls/build_rules:xls_toolchains.bzl",
@@ -140,6 +144,9 @@ def _xls_ir_jit_wrapper_impl(ctx):
         fail("Header filename must contain the '%s' extension." %
              _H_FILE_EXTENSION)
 
+    if ctx.attr.top != "":
+        jit_wrapper_flags.add("--function", ctx.attr.top)
+
     # validate basename
     if source_basename != header_basename:
         fail("The basename of the source and header files do not match.")
@@ -154,7 +161,7 @@ def _xls_ir_jit_wrapper_impl(ctx):
     jit_wrapper_flags.add("--output_dir", cc_file.dirname)
 
     # genfiles directory
-    jit_wrapper_flags.add("--genfiles_dir", ctx.genfiles_dir.path)
+    jit_wrapper_flags.add("--genfiles_dir", ctx.bin_dir.path + "/")
 
     # function_type
     jit_wrapper_flags.add("--function_type", ctx.attr.wrapper_type)
@@ -221,6 +228,7 @@ Examples:
     implementation = _xls_ir_jit_wrapper_impl,
     attrs = dicts.add(
         xls_ir_common_attrs,
+        xls_ir_top_attrs,
         _xls_ir_jit_wrapper_attrs,
         CONFIG["xls_outs_attrs"],
         xls_toolchain_attrs,
@@ -248,6 +256,7 @@ _no_aot_info = rule(
 def xls_ir_jit_wrapper_macro(
         name,
         src,
+        top,
         source_file,
         header_file,
         wrapper_type,
@@ -265,6 +274,7 @@ def xls_ir_jit_wrapper_macro(
     Args:
       name: The name of the rule.
       src: The IR file. See 'src' attribute from the 'xls_ir_jit_wrapper' rule.
+      top: Name of the function/block/proc to wrap.
       source_file: The generated source file. See 'source_file' attribute from
         the 'xls_ir_jit_wrapper' rule.
       header_file: The generated header file. See 'header_file' attribute from
@@ -284,6 +294,7 @@ def xls_ir_jit_wrapper_macro(
     # Type check input
     string_type_check("name", name)
     string_type_check("src", src)
+    string_type_check("top", top)
     string_type_check("source_file", source_file)
     string_type_check("header_file", header_file)
     string_type_check("wrapper_type", wrapper_type)
@@ -295,6 +306,7 @@ def xls_ir_jit_wrapper_macro(
     xls_ir_jit_wrapper(
         name = name,
         src = src,
+        top = top,
         source_file = source_file,
         header_file = header_file,
         aot_info = aot_info,
@@ -316,17 +328,24 @@ PROC_WRAPPER_TYPE = "PROC"
 BLOCK_WRAPPER_TYPE = "BLOCK"
 
 _BASE_JIT_WRAPPER_DEPS = {
-    FUNCTION_WRAPPER_TYPE: "//xls/jit:function_base_jit_wrapper",
-    PROC_WRAPPER_TYPE: "//xls/jit:proc_base_jit_wrapper",
+    FUNCTION_WRAPPER_TYPE: ["//xls/jit:function_base_jit_wrapper"],
+    PROC_WRAPPER_TYPE: [
+        "//xls/jit:proc_base_jit_wrapper",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+    BLOCK_WRAPPER_TYPE: [
+        "//xls/jit:block_base_jit_wrapper",
+        "//xls/ir:bits",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
 }
 
-_AOT_SUPPORTED_WRAPPERS = [FUNCTION_WRAPPER_TYPE, PROC_WRAPPER_TYPE]
-
 def cc_xls_ir_jit_wrapper(
         name,
         src,
         jit_wrapper_args = {},
         wrapper_type = FUNCTION_WRAPPER_TYPE,
+        top = "",
         **kwargs):
     """Invokes the JIT wrapper generator and compiles the result as a cc_library.
 
@@ -344,13 +363,13 @@ def cc_xls_ir_jit_wrapper(
                     FUNCTION_WRAPPER_TYPE, BLOCK_WRAPPER_TYPE, or
                     PROC_WRAPPER_TYPE symbols. Defaults to FUNCTION_WRAPPER_TYPE
                     for compatibility.
+      top: Name of the top function/proc/block.
       **kwargs: Keyword arguments. Named arguments.
     """
 
-    # TODO(allight): We should add top as an argument here. With the new
-    # jit-wrapper architecture it would be simple to support.
     dictionary_type_check("jit_wrapper_args", jit_wrapper_args)
     string_type_check("src", src)
+    string_type_check("top", src)
 
     # Validate arguments of macro
     if kwargs.get("source_file"):
@@ -363,31 +382,25 @@ def cc_xls_ir_jit_wrapper(
     if wrapper_type not in (FUNCTION_WRAPPER_TYPE, BLOCK_WRAPPER_TYPE, PROC_WRAPPER_TYPE):
         fail(("Cannot set 'wrapper_type' to %s. It must be one of BLOCK_WRAPPER_TYPE, " +
               "FUNCTION_WRAPPER_TYPE, or PROC_WRAPPER_TYPE") % wrapper_type)
-    if wrapper_type == BLOCK_WRAPPER_TYPE:
-        fail("Block jit-wrapper not yet supported!")
 
     source_filename = name + _CC_FILE_EXTENSION
     header_filename = name + _H_FILE_EXTENSION
 
     extra_lib_deps = []
-    if wrapper_type in _AOT_SUPPORTED_WRAPPERS:
-        xls_aot_generate(
-            name = name + "_aot_code_for_wrapper",
-            src = src,
-            with_msan = XLS_IS_MSAN_BUILD,
-            # The XLS AOT compiler does not currently support cross-compilation.
-        )
-        aot_info_target = ":" + name + "_aot_code_for_wrapper"
-        extra_lib_deps.append(aot_info_target)
-    else:
-        # Simplify the xls_ir_jit_wrapper_macro by making sure it always gets an AotCompileInfo
-        _no_aot_info(name = name + "_empty_aot_info")
-        aot_info_target = ":" + name + "_empty_aot_info"
-        # Since this doesn't define any actual AOT code we don't need to add anything to the deps.
+    xls_aot_generate(
+        name = name + "_aot_code_for_wrapper",
+        src = src,
+        top = top,
+        with_msan = XLS_IS_MSAN_BUILD,
+        # The XLS AOT compiler does not currently support cross-compilation.
+    )
+    aot_info_target = ":" + name + "_aot_code_for_wrapper"
+    extra_lib_deps.append(aot_info_target)
 
     xls_ir_jit_wrapper_macro(
         name = "__" + name + "_xls_ir_jit_wrapper",
         src = src,
+        top = top,
         jit_wrapper_args = jit_wrapper_args,
         aot_info = aot_info_target,
         wrapper_type = wrapper_type,
@@ -400,17 +413,12 @@ def cc_xls_ir_jit_wrapper(
         name = name,
         srcs = [":" + source_filename],
         hdrs = [":" + header_filename],
-        deps = extra_lib_deps + [
-            _BASE_JIT_WRAPPER_DEPS[wrapper_type],
-            "@com_google_absl//absl/container:flat_hash_map",
+        deps = extra_lib_deps +
+               _BASE_JIT_WRAPPER_DEPS[wrapper_type] + [
             "@com_google_absl//absl/status",
             "@com_google_absl//absl/status:statusor",
             "//xls/common/status:status_macros",
             "//xls/interpreter:evaluator_options",
-            "//xls/jit:function_jit",
-            "//xls/public:function_builder",
-            "//xls/public:ir",
-            "//xls/public:ir_parser",
             "//xls/public:value",
         ],
         **kwargs

diff --git a/xls/jit/BUILD b/xls/jit/BUILD
@@ -28,6 +28,8 @@ load(
     "xls_dslx_library",
     "xls_dslx_opt_ir",
     "xls_ir_cc_library",
+    "xls_ir_opt_ir",
+    "xls_ir_verilog",
 )
 load(
     "//xls/build_rules:xls_internal_aot_rules.bzl",
@@ -91,6 +93,7 @@ cc_binary(
     srcs = ["aot_compiler_main.cc"],
     deps = [
         ":aot_entrypoint_cc_proto",
+        ":block_jit",
         ":function_base_jit",
         ":function_jit",
         ":jit_proc_runtime",
@@ -102,6 +105,7 @@ cc_binary(
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/ir",
+        "//xls/ir:block_elaboration",
         "//xls/ir:ir_parser",
         "//xls/ir:type",
         "@com_google_absl//absl/algorithm:container",
@@ -232,6 +236,30 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "block_base_jit_wrapper",
+    srcs = ["block_base_jit_wrapper.cc"],
+    hdrs = ["block_base_jit_wrapper.h"],
+    # Allow jit-wrapper users to see this.
+    visibility = ["//xls:xls_users"],
+    deps = [
+        ":block_jit",
+        ":function_base_jit",
+        "//xls/common/status:ret_check",
+        "//xls/common/status:status_macros",
+        "//xls/interpreter:evaluator_options",
+        "//xls/ir",
+        "//xls/ir:value",
+        "//xls/ir:value_utils",
+        "//xls/public:ir_parser",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
 cc_library(
     name = "proc_base_jit_wrapper",
     hdrs = ["proc_base_jit_wrapper.h"],
@@ -261,6 +289,8 @@ py_binary(
     name = "jit_wrapper_generator_main",
     srcs = ["jit_wrapper_generator_main.py"],
     data = [
+        ":jit_block_wrapper_cc.tmpl",
+        ":jit_block_wrapper_h.tmpl",
         ":jit_function_wrapper_cc.tmpl",
         ":jit_function_wrapper_h.tmpl",
         ":jit_proc_wrapper_cc.tmpl",
@@ -285,6 +315,7 @@ cc_test(
     srcs = ["jit_wrapper_test.cc"],
     deps = [
         ":compound_type_jit_wrapper",
+        ":multi_func_block_wrapper",
         "//xls/common:xls_gunit_main",
         "//xls/common/status:matchers",
         "//xls/dslx/stdlib:float32_mul_jit_wrapper",
@@ -709,6 +740,8 @@ cc_library(
     srcs = ["block_jit.cc"],
     hdrs = ["block_jit.h"],
     deps = [
+        ":aot_compiler",
+        ":aot_entrypoint_cc_proto",
         ":function_base_jit",
         ":jit_buffer",
         ":jit_callbacks",
@@ -740,6 +773,8 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:ir_headers",
     ],
 )
 
@@ -1005,7 +1040,10 @@ build_test(
 proto_library(
     name = "aot_entrypoint_proto",
     srcs = ["aot_entrypoint.proto"],
-    deps = [":type_layout_proto"],
+    deps = [
+        ":type_layout_proto",
+        "//xls/ir:xls_type_proto",
+    ],
 )
 
 cc_proto_library(
@@ -1062,6 +1100,46 @@ xls_dslx_ir(
     library = ":multi_function_dslx",
 )
 
+xls_ir_opt_ir(
+    name = "multi_func_opt_ir",
+    src = ":multi_function",
+)
+
+xls_ir_verilog(
+    name = "multi_func_sv",
+    src = ":multi_func_opt_ir",
+    codegen_args = {
+        "module_name": "multi_func_top",
+        "delay_model": "unit",
+        "generator": "combinational",
+        "reset": "rst",
+        "reset_data_path": "true",
+        "reset_active_low": "false",
+        "reset_asynchronous": "false",
+        "flop_inputs": "false",
+        "flop_single_value_channels": "false",
+        "flop_outputs": "false",
+        "add_idle_output": "false",
+        "multi_proc": "true",
+        "streaming_channel_data_suffix": "_data",
+        "streaming_channel_ready_suffix": "_ready",
+        "streaming_channel_valid_suffix": "_valid",
+        "use_system_verilog": "true",
+    },
+    verilog_file = "multi_func.sv",
+)
+
+cc_xls_ir_jit_wrapper(
+    name = "multi_func_block_wrapper",
+    src = ":multi_func.block.ir",
+    jit_wrapper_args = {
+        "class_name": "MultiFuncBlockJit",
+        "namespace": "something::cool",
+    },
+    top = "multi_func_top",
+    wrapper_type = "BLOCK",
+)
+
 xls_aot_generate(
     name = "multi_function_aot",
     testonly = True,