Update on "[executorch][serialization] Data serialization interface"

Introduce data serialization interface.

Differential Revision: [D65947145](https://our.internmc.facebook.com/intern/diff/D65947145/)

[ghstack-poisoned]

Author: lucylq

backends/apple/coreml/runtime/test/export_stateful_model.py

@@ -47,7 +47,7 @@ def main() -> None:
         torch.randn((1, embedding_dim)),
         torch.tensor([0]),
     )
-    exported_model = export(model, example_inputs)
+    exported_model = export(model, example_inputs, strict=True)
     edge_program_manager = exir.to_edge(exported_model)
     compile_specs = CoreMLBackend.generate_compile_specs(
         compute_precision=ct.precision.FLOAT16,

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -16,7 +16,6 @@
 
 
 class TestCoreMLPartitioner(unittest.TestCase):
-
     # TODO(T182928844): Delegate dim order op to backend.
     edge_compile_config = executorch.exir.EdgeCompileConfig(_skip_dim_order=True)
 
@@ -34,7 +33,7 @@ def forward(self, a, x, b):
         model.eval()
 
         example_inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
 
         edge_program_manager = executorch.exir.to_edge(
             exir_program_aten, compile_config=self.edge_compile_config
@@ -61,7 +60,7 @@ def test_vit_skip_conv(self):
         model.eval()
 
         example_inputs = (torch.randn(1, 3, 224, 224),)
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
         edge_program_manager = executorch.exir.to_edge(
             exir_program_aten, compile_config=self.edge_compile_config
         )
@@ -106,7 +105,7 @@ def forward(self, q, k_val, input_pos):
         k_val = torch.randn((1, embedding_dim))
         input_pos = torch.tensor([0])
         example_inputs = (q, k_val, input_pos)
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
 
         compile_specs = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18

backends/apple/mps/test/test_mps_utils.py

@@ -247,10 +247,7 @@ def lower_module_and_test_output(
             )
 
             executorch_program = to_edge(
-                export(
-                    delegated_program,
-                    sample_inputs,
-                ),
+                export(delegated_program, sample_inputs, strict=True),
                 compile_config=exir.EdgeCompileConfig(
                     _check_ir_validity=False,
                     _skip_dim_order=True,  # TODO(T182928844): Delegate dim order op to backend.

backends/cadence/aot/compiler.py

@@ -176,7 +176,7 @@ def export_program(
     torch._C._set_mkldnn_enabled(False)
 
     # else: capture the model and return it.
-    expo_program = export(model, inputs)
+    expo_program = export(model, inputs, strict=True)
 
     if dump_graphs:
         logging.info("Exported graph:")

backends/example/test_example_delegate.py

@@ -60,7 +60,7 @@ def get_example_inputs():
 
         quantized_gm = m
         exported_program = to_edge(
-            export(quantized_gm, copy.deepcopy(example_inputs)),
+            export(quantized_gm, copy.deepcopy(example_inputs), strict=True),
             compile_config=EDGE_COMPILE_CONFIG,
         )
 
@@ -92,7 +92,7 @@ def test_delegate_mobilenet_v2(self):
 
         quantized_gm = m
         exported_program = to_edge(
-            export(quantized_gm, copy.deepcopy(example_inputs)),
+            export(quantized_gm, copy.deepcopy(example_inputs), strict=True),
             compile_config=EDGE_COMPILE_CONFIG,
         )
 

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1617,7 +1617,7 @@ def test_qnn_backend_multi_contexts_composite(self):
         )
         sample_input = module.get_random_input()
         edge_prog = to_edge(
-            torch.export.export(module, sample_input),
+            torch.export.export(module, sample_input, strict=True),
         )
         update_spill_fill_size(edge_prog.exported_program())
         exec_prog = edge_prog.to_executorch()
@@ -1957,7 +1957,7 @@ def calibrator(gm):
         self.assertEqual(len(exported_progs), 1)
         # lower all graph again, the skipped operators will be left in CPU
         exec_prog = to_edge(
-            torch.export.export(graph_module, sample_input),
+            torch.export.export(graph_module, sample_input, strict=True),
         ).to_executorch()
         self.verify_output(module, sample_input, exec_prog)
 
@@ -2004,7 +2004,7 @@ def calibrator(gm):
         self.assertEqual(len(exported_progs), 2)
         # lower all graph again, the skipped operators will be left in CPU
         exec_prog = exec_prog = to_edge(
-            torch.export.export(graph_module, sample_input),
+            torch.export.export(graph_module, sample_input, strict=True),
         ).to_executorch()
         self.verify_output(module, sample_input, exec_prog)
 
@@ -2041,7 +2041,7 @@ def calibrator(gm):
         self.assertEqual(len(exported_progs), 5)
         # lower all graph again, the skipped operators will be delegated with fp16
         exec_prog = to_edge(
-            torch.export.export(graph_module, sample_input),
+            torch.export.export(graph_module, sample_input, strict=True),
         ).to_executorch()
         self.verify_output(module, sample_input, exec_prog)
 
@@ -2086,7 +2086,7 @@ def test_qnn_backend_multi_contexts_composite(self):
         )
         sample_input = module.get_random_input()
         edge_prog = to_edge(
-            torch.export.export(module, sample_input),
+            torch.export.export(module, sample_input, strict=True),
         )
         update_spill_fill_size(edge_prog.exported_program())
         exec_prog = edge_prog.to_executorch()
@@ -2721,7 +2721,6 @@ def test_ssd300_vgg16(self):
 
 
 class TestExampleQaihubScript(TestQNN):
-
     def required_envs(self, conditions=None) -> bool:
         conditions = [] if conditions is None else conditions
         return all(

backends/qualcomm/tests/utils.py

@@ -385,7 +385,7 @@ def get_qdq_module(
         custom_quant_annotations: Tuple[Callable] = (),
         quant_dtype: QuantDtype = QuantDtype.use_8a8w,
     ) -> torch.fx.GraphModule:
-        m = torch.export.export(module, inputs).module()
+        m = torch.export.export(module, inputs, strict=True).module()
 
         quantizer = QnnQuantizer()
         quantizer.add_custom_quant_annotations(custom_quant_annotations)

backends/qualcomm/utils/utils.py

@@ -337,7 +337,7 @@ def capture_program(
     inputs: Tuple[torch.Tensor],
     custom_pass_config: FrozenSet[str] = frozenset(),
 ) -> exir.ExirExportedProgram:
-    ep = torch.export.export(module, inputs)
+    ep = torch.export.export(module, inputs, strict=True)
     decomposed_ep = ep.run_decompositions(get_decomp_table())
     # We choose call_operator by target in ConvertBinaryOpsWithScalar
     # because it is the same source_fn_stack for MultiheadAttention
@@ -551,7 +551,7 @@ def prepare_subgm(subgm, subgm_name):
 
     fp_node_id_set = fp_node_id_set if fp_node_id_set is not None else set()
     fp_node_op_set = fp_node_op_set if fp_node_op_set is not None else set()
-    graph_module = torch.export.export(nn_module, sample_input).module()
+    graph_module = torch.export.export(nn_module, sample_input, strict=True).module()
     # define node support type
     capability_partitioner = CapabilityBasedPartitioner(
         graph_module,
@@ -664,7 +664,7 @@ def forward(self, *inputs):
                 ).default(inputs)
 
         model = Model()
-        prog = torch.export.export(model, tuple(inputs.values()))
+        prog = torch.export.export(model, tuple(inputs.values()), strict=True)
         # bookkeeping for variables' life cycle
         return {
             "custom_op": custom_op,

backends/vulkan/runtime/graph/ops/glsl/copy_channel_offset.glsl

@@ -18,17 +18,16 @@ ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "existing_out", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 
-${layout_declare_ubo(B, "ivec4", "out_sizes")}
-${layout_declare_ubo(B, "ivec4", "in_sizes")}
-
-layout(set = 0, binding = 5) uniform PRECISION restrict CopyArgs {
+layout(push_constant) uniform restrict Block {
+  ivec4 out_sizes;
+  ivec4 in_sizes;
   // Operates on (x, y, z) logical extents.
-  ivec3 range;
+  // channel_range is stored in range.w
+  ivec4 range;
   // Analogus to range variable in copy. It defines the # of channel being
   // copied.
-  int channel_range;
-  ivec3 dst_offset;
-  int dst_channel_offset;
+  // dst channel offset is stored in dst_offset.w
+  ivec4 dst_offset;
   int src_channel_offset;
 };
 
@@ -47,11 +46,11 @@ void main() {
   // Note: Unlike other shaders, the range is often not equal to the destination
   // texture extent.
   const ivec3 lpos = ivec3(gl_GlobalInvocationID);
-  if (any(greaterThanEqual(lpos, range))) {
+  if (any(greaterThanEqual(lpos, range.xyz))) {
     return;
   }
 
-  const ivec3 out_lpos = lpos + dst_offset;
+  const ivec3 out_lpos = lpos + dst_offset.xyz;
 
   const ivec4 out_tidx = lpos_to_tidx(out_lpos, out_sizes, out_axis_map.w, packed_dim);
 
@@ -61,12 +60,12 @@ void main() {
   ivec4 in_tidx = out_tidx;
   for (int i=0; i<4; i++) {
 
-    in_tidx[packed_dim] = out_tidx[packed_dim] - dst_channel_offset + i;
+    in_tidx[packed_dim] = out_tidx[packed_dim] - dst_offset.w + i;
 
     // Handle the partial update for begining of channel in an existing tensor.
     // If the source channel index is below zero or exceeds the range, we skip
     // updating the element to avoid overwriting existing data.
-    if ((in_tidx[packed_dim] < 0) || (in_tidx[packed_dim] >= channel_range)) {
+    if ((in_tidx[packed_dim] < 0) || (in_tidx[packed_dim] >= range.w)) {
       continue;
     }
 

backends/vulkan/runtime/graph/ops/glsl/copy_offset.glsl

@@ -17,7 +17,11 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 
-${layout_declare_ubo(B, "ivec3", "range", "ivec3", "src_offset", "ivec3", "dst_offset")}
+layout(push_constant) uniform restrict Block {
+  ivec3 range;
+  ivec3 src_offset;
+  ivec3 dst_offset;
+};
 
 #include "indexing_utils.h"
 

backends/vulkan/runtime/graph/ops/impl/Copy.cpp

@@ -33,16 +33,6 @@ void add_copy_offset_node(
   add_dtype_suffix(kernel_name, *t_out);
   add_storage_type_suffix(kernel_name, *t_out);
 
-  const struct Block final {
-    alignas(16) ivec3 range;
-    alignas(16) ivec3 src_offset;
-    alignas(16) ivec3 dst_offset;
-  } offset_params{
-      range,
-      src_offset,
-      dst_offset,
-  };
-
   auto shader = VK_KERNEL_FROM_STR(kernel_name);
 
   graph.execute_nodes().emplace_back(new DispatchNode(
@@ -56,11 +46,18 @@ void add_copy_offset_node(
           {in, vkapi::kRead},
       },
       // Parameter buffers
-      {
-          graph.create_params_buffer(offset_params),
-      },
+      {},
       // Specialization Constants
-      {graph.hashed_layout_of(out), graph.hashed_layout_of(in)}));
+      {graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
+      nullptr,
+      {},
+      {
+          PushConstantDataInfo(&range, sizeof(range), sizeof(utils::ivec4)),
+          PushConstantDataInfo(
+              &src_offset, sizeof(src_offset), sizeof(utils::ivec4)),
+          PushConstantDataInfo(
+              &dst_offset, sizeof(dst_offset), sizeof(utils::ivec4)),
+      }));
 }
 
 void add_copy_channel_offset_node(
@@ -128,28 +125,23 @@ void add_copy_channel_offset_node(
     // The shader combines the global invocation id and the dst_offset to get
     // the actual coordinate.
 
-    ivec3 dst_offset{
+    const ivec3 dst_offset{
         0, 0, dst_first_z + batch_idx * utils::div_up_4(out_channels)};
 
-    uvec3 global_size{
+    const uvec3 global_size{
         utils::safe_downcast<uint32_t>(dim_at<kWidth4D>(in_sizes)),
         utils::safe_downcast<uint32_t>(dim_at<kHeight4D>(in_sizes)),
         utils::safe_downcast<uint32_t>(dst_last_z - dst_first_z + 1)};
-    uvec3 local_size = graph.create_local_wg_size(global_size);
-
-    const struct Block final {
-      ivec3 range;
-      int32_t channel_range;
-      ivec3 dst_offset;
-      int32_t dst_channel_offset;
-      int32_t src_channel_offset;
-    } channel_offset_params{
-        utils::make_ivec3(global_size),
-        channel_range,
-        dst_offset,
-        dst_channel_offset,
-        src_channel_offset,
-    };
+    const uvec3 local_size = graph.create_local_wg_size(global_size);
+
+    const utils::ivec4 range_params = {
+        static_cast<int>(global_size[0]),
+        static_cast<int>(global_size[1]),
+        static_cast<int>(global_size[2]),
+        channel_range};
+
+    const utils::ivec4 offset_params = {
+        dst_offset[0], dst_offset[1], dst_offset[2], dst_channel_offset};
 
     auto shader = VK_KERNEL_FROM_STR(kernel_name);
 
@@ -165,13 +157,17 @@ void add_copy_channel_offset_node(
             {in, vkapi::MemoryAccessType::READ},
         },
         // Parameter buffers
-        {
-            t_out->sizes_ubo(),
-            t_in->sizes_ubo(),
-            graph.create_params_buffer(channel_offset_params),
-        },
+        {},
         // Specialization Constants
-        {graph.hashed_layout_of(out), graph.hashed_layout_of(in)}));
+        {graph.hashed_layout_of(out), graph.hashed_layout_of(in)},
+        nullptr,
+        {},
+        {graph.sizes_pc_of(out),
+         graph.sizes_pc_of(in),
+         PushConstantDataInfo(&range_params, sizeof(range_params)),
+         PushConstantDataInfo(&offset_params, sizeof(offset_params)),
+         PushConstantDataInfo(
+             &src_channel_offset, sizeof(src_channel_offset))}));
   }
 }
 

backends/vulkan/targets.bzl

@@ -1,4 +1,4 @@
-load("@fbcode//target_determinator/macros:ci.bzl", "ci")
+load("@fbsource//tools/target_determinator/macros:ci.bzl", "ci")
 load("@fbcode_macros//build_defs:native_rules.bzl", "buck_genrule")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID", "CXX", "FBCODE")

backends/vulkan/test/compute_api_tests.bzl

@@ -1,4 +1,4 @@
-load("@fbcode//target_determinator/macros:ci.bzl", "ci")
+load("@fbsource//tools/target_determinator/macros:ci.bzl", "ci")
 load("@fbsource//tools/build_defs:fb_xplat_cxx_binary.bzl", "fb_xplat_cxx_binary")
 load("@fbsource//tools/build_defs:fb_xplat_cxx_test.bzl", "fb_xplat_cxx_test")
 load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID", "MACOSX", "CXX")

backends/vulkan/test/test_vulkan_delegate.py

@@ -112,7 +112,7 @@ def run_test():
             model(*sample_inputs)
 
             program: ExportedProgram = export(
-                model, sample_inputs, dynamic_shapes=dynamic_shapes
+                model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True
             )
 
             edge_program = to_edge_transform_and_lower(

backends/xnnpack/partition/graphs/sdpa.py

@@ -76,6 +76,7 @@ def forward(
                         v,
                         mask,
                     ),
+                    strict=True,
                 ),
                 compile_config=get_xnnpack_edge_compile_config(),
             )

backends/xnnpack/test/tester/tester.py

@@ -194,7 +194,7 @@ def run(
         inputs: Tuple[torch.Tensor],
     ) -> None:
         self.exported_program = export(
-            artifact, inputs, dynamic_shapes=self.dynamic_shapes
+            artifact, inputs, dynamic_shapes=self.dynamic_shapes, strict=True
         )
 
     @property

build/packaging/smoke_test.py

@@ -65,7 +65,7 @@ def export_linear_model() -> bytes:
 
     # Export the pytorch model and process for ExecuTorch.
     print("Exporting program...")
-    exported_program = export(LinearModel(), example_inputs)
+    exported_program = export(LinearModel(), example_inputs, strict=True)
     print("Lowering to edge...")
     edge_program = to_edge(exported_program)
     print("Creating ExecuTorch program...")

devtools/backend_debug/tests/test_delegation_info.py

@@ -31,7 +31,7 @@ def forward(self, a, x, b):
 
         m = Model()
         inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-        edge = to_edge(torch.export.export(m, inputs)).to_backend(
+        edge = to_edge(torch.export.export(m, inputs, strict=True)).to_backend(
             AddMulPartitionerDemo()
         )
         delegation_info = get_delegation_info(edge.exported_program().graph_module)