[torch-frontend] use min(target_version, current_version) as target_v… (

#460) …ersion
bytedance · Oct 12, 2024 · c0e688b · c0e688b
1 parent 5b9a49e
commit c0e688b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/frontends/torch-frontend/torch-frontend/python/test/test_stablehlo_bytecode.py b/frontends/torch-frontend/torch-frontend/python/test/test_stablehlo_bytecode.py
@@ -9,6 +9,9 @@ def serialize_helper(module, inputs):
     stablehlo_bytecode = compile(module, inputs, "stablehlo+0.16.2")
     deserialize_str = deserialize_portable_artifact(stablehlo_bytecode)
     print(deserialize_str)
+    stablehlo_bytecode = compile(module, inputs, "stablehlo+10.0.0")
+    deserialize_str = deserialize_portable_artifact(stablehlo_bytecode)
+    print(deserialize_str)
 
 # ==============================================================================
 class SoftmaxModule(torch.nn.Module):

diff --git a/frontends/torch-frontend/torch-frontend/python/torch_frontend/compile.py b/frontends/torch-frontend/torch-frontend/python/torch_frontend/compile.py
@@ -11,7 +11,7 @@
 
 from . import ir
 from .passmanager import PassManager
-from ._mlir_libs._stablehlo import serialize_portable_artifact
+from ._mlir_libs._stablehlo import serialize_portable_artifact, get_current_version
 
 from .extra_shape_fn import byteir_extra_library
 
@@ -259,8 +259,13 @@ def compile(
     ############################################
     # serialize stablehlo to target version
     ############################################
+    from packaging import version
+    target_version = version.Version(output_type.split("+")[1])
+    current_version = version.Version(get_current_version())
+    if target_version > current_version:
+        target_version = current_version
     return serialize_portable_artifact(
-        module.operation.get_asm(), output_type.split("+")[1]
+        module.operation.get_asm(), str(target_version)
     )
 
 
@@ -380,6 +385,11 @@ def compile_dynamo_model(
     ############################################
     # serialize stablehlo to target version
     ############################################
+    from packaging import version
+    target_version = version.Version(output_type.split("+")[1])
+    current_version = version.Version(get_current_version())
+    if target_version > current_version:
+        target_version = current_version
     return serialize_portable_artifact(
-        module.operation.get_asm(), output_type.split("+")[1]
+        module.operation.get_asm(), str(target_version)
     )