[export] Sync aoti schema to schema.py (#148017)

zhxchen17 · pytorchmergebot · commit 915b9c80ab59 · 2025-02-27T21:46:11.000Z
Summary: Synchronizing internal AOTI schema to OSS schema.py Test Plan: CI Differential Revision: D70271151 Pull Request resolved: #148017 Approved by: https://github.com/yiming0416
diff --git a/torch/_export/serde/export_schema.thrift b/torch/_export/serde/export_schema.thrift
@@ -1,5 +1,5 @@
 // @generated by update_schema.py
-// checksum<<735727a20b699856df15a522b7e0075dae34d441a021be640f447eb245b426c9>>
+// checksum<<f36968728ea96d9629b7c5269f5303e5cf23fba341d0221cb364aaf571b94dd6>>
 
 namespace py3 torch._export
 namespace cpp2 torch._export.schema
@@ -341,3 +341,21 @@ struct Model {
   60: map<string, string> deviceAllocationMap;
   70: map<string, string> constantPaths;
 }
+
+struct AOTInductorModelPickleData {
+  1: string library_basename;
+  2: list<string> input_names;
+  3: list<string> output_names;
+  4: optional i64 floating_point_input_dtype;
+  5: optional i64 floating_point_output_dtype;
+  6: optional bool aot_inductor_model_is_cpu;
+}
+
+struct ExternKernelNode {
+  10: string name;
+  20: Node node;
+}
+
+struct ExternKernelNodes {
+  10: list<ExternKernelNode> nodes;
+}
diff --git a/torch/_export/serde/schema.py b/torch/_export/serde/schema.py
@@ -8,7 +8,7 @@
 from torch._export.serde.union import _Union
 
 # NOTE: Please update this value if any modifications are made to the schema
-SCHEMA_VERSION = (8, 6)
+SCHEMA_VERSION = (8, 7)
 TREESPEC_VERSION = 1
 
 
@@ -447,3 +447,40 @@ class Model:
     # key is the FQN of constant in exported program (constant tensor or torchbind objs)
     # value is the archive path of serialized constants
     constantPaths: Annotated[dict[str, str], 70]
+
+#
+# The structure is used to serialize instances of AOTInductorModel to pass
+# them from the publishing pipeline to the predictor.
+#
+# All new fields should be marked as optional.
+#
+@dataclass
+class AOTInductorModelPickleData:
+    # Base name of an associated .so AOTInductor library. Typically looks like:
+    # "abc.so".
+    library_basename: Annotated[str, 1]
+
+    # AOTInductor engine input names.
+    input_names: Annotated[list[str], 2]
+
+    # AOTInductor engine output names.
+    output_names: Annotated[list[str], 3]
+
+    # These fields tell whether floating point inputs/outputs should be converted to
+    # a certain type. If None, the dtypes that the AOTInductor engine inferred from the sample
+    # inputs are used.
+    floating_point_input_dtype: Annotated[Optional[int], 4] = None
+    floating_point_output_dtype: Annotated[Optional[int], 5] = None
+
+    # Whether AOTInductor runtime is for CPU.
+    aot_inductor_model_is_cpu: Annotated[Optional[bool], 6] = None
+
+@dataclass
+class ExternKernelNode:
+    # name is not the unique identifier of the node
+    name: Annotated[str, 10]
+    node: Annotated[Node, 20]
+
+@dataclass
+class ExternKernelNodes:
+    nodes: Annotated[list[ExternKernelNode], 10]
diff --git a/torch/_export/serde/schema.yaml b/torch/_export/serde/schema.yaml
@@ -1,5 +1,23 @@
 # @generated by update_schema.py
-# checksum<<4710fa9fbec49ef1c60365f9eccade2b963e259d6aafc42ce1b32a4f2e20a858>>
+# checksum<<31c433c768b3f1bb61a5e8f4ceffc40c857bd80cf4fa0fc33fd03fa5ebb6c4d8>>
+AOTInductorModelPickleData:
+  kind: struct
+  fields:
+    library_basename:
+      type: str
+    input_names:
+      type: List[str]
+    output_names:
+      type: List[str]
+    floating_point_input_dtype:
+      type: Optional[int]
+      default: None
+    floating_point_output_dtype:
+      type: Optional[int]
+      default: None
+    aot_inductor_model_is_cpu:
+      type: Optional[bool]
+      default: None
 Argument:
   kind: union
   fields:
@@ -111,6 +129,18 @@ ExportedProgram:
     torch_version:
       type: str
       default: <=2.4
+ExternKernelNode:
+  kind: struct
+  fields:
+    name:
+      type: str
+    node:
+      type: Node
+ExternKernelNodes:
+  kind: struct
+  fields:
+    nodes:
+      type: List[ExternKernelNode]
 GradientToParameterSpec:
   kind: struct
   fields:
@@ -500,5 +530,5 @@ UserOutputSpec:
       type: Argument
 SCHEMA_VERSION:
 - 8
-- 6
+- 7
 TREESPEC_VERSION: 1
diff --git a/torch/csrc/utils/generated_serialization_types.h b/torch/csrc/utils/generated_serialization_types.h
