From 661258a6eb19902697deda5fb7eb63ee8e54ef4a Mon Sep 17 00:00:00 2001
From: Xuanda Yang <th3charlie@gmail.com>
Date: Mon, 6 Nov 2023 15:46:16 -0800
Subject: [PATCH] add missing serialization of Dim::partition_policy

---
 src/Deserialization.cpp | 2 ++
 src/Serialization.cpp   | 3 ++-
 src/halide_ir.fbs       | 1 +
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp
index eda8ad93338b..bb19cf82c9aa 100644
--- a/src/Deserialization.cpp
+++ b/src/Deserialization.cpp
@@ -1110,11 +1110,13 @@ Dim Deserializer::deserialize_dim(const Serialize::Dim *dim) {
     const auto for_type = deserialize_for_type(dim->for_type());
     const auto device_api = deserialize_device_api(dim->device_api());
     const auto dim_type = deserialize_dim_type(dim->dim_type());
+    const auto partition_policy = deserialize_partition(dim->partition_policy());
     auto hl_dim = Dim();
     hl_dim.var = var;
     hl_dim.for_type = for_type;
     hl_dim.device_api = device_api;
     hl_dim.dim_type = dim_type;
+    hl_dim.partition_policy = partition_policy;
     return hl_dim;
 }
 
diff --git a/src/Serialization.cpp b/src/Serialization.cpp
index 038d5a1323e0..c85eaa15e1aa 100644
--- a/src/Serialization.cpp
+++ b/src/Serialization.cpp
@@ -1227,7 +1227,8 @@ Offset<Serialize::Dim> Serializer::serialize_dim(FlatBufferBuilder &builder, con
     const auto for_type_serialized = serialize_for_type(dim.for_type);
     const auto device_api_serialized = serialize_device_api(dim.device_api);
     const auto dim_type_serialized = serialize_dim_type(dim.dim_type);
-    return Serialize::CreateDim(builder, var_serialized, for_type_serialized, device_api_serialized, dim_type_serialized);
+    const auto partition_policy_serialized = serialize_partition(dim.partition_policy);
+    return Serialize::CreateDim(builder, var_serialized, for_type_serialized, device_api_serialized, dim_type_serialized, partition_policy_serialized);
 }
 
 Offset<Serialize::FuseLoopLevel> Serializer::serialize_fuse_loop_level(FlatBufferBuilder &builder, const FuseLoopLevel &fuse_loop_level) {
diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs
index f0ad94d8bdf6..f3d27e83a62a 100644
--- a/src/halide_ir.fbs
+++ b/src/halide_ir.fbs
@@ -558,6 +558,7 @@ table Dim {
     for_type: ForType;
     device_api: DeviceAPI;
     dim_type: DimType;
+    partition_policy: Partition;
 }
 
 enum LoopAlignStrategy: ubyte {