use dict.setdefault in with_fx_graph

siliconflow · Jan 14, 2024 · 3a745a6 · 3a745a6
1 parent fad4dd4
commit 3a745a6
Showing 1 changed file with 14 additions and 14 deletions.
diff --git a/src/onediff/infer_compiler/with_fx_graph.py b/src/onediff/infer_compiler/with_fx_graph.py
@@ -20,20 +20,20 @@ def fx_node_tranform(gm):
     if not enable_graph:
         oneflow_fn = of_gm.forward
     else:
-        os.environ["ONEFLOW_MLIR_CSE"] = "1"
-        os.environ["ONEFLOW_MLIR_ENABLE_INFERENCE_OPTIMIZATION"] = "1"
-        os.environ["ONEFLOW_MLIR_ENABLE_ROUND_TRIP"] = "1"
-        os.environ["ONEFLOW_MLIR_FUSE_FORWARD_OPS"] = "1"
-        os.environ["ONEFLOW_MLIR_FUSE_OPS_WITH_BACKWARD_IMPL"] = "1"
-        os.environ["ONEFLOW_MLIR_GROUP_MATMUL"] = "1"
-        os.environ["ONEFLOW_MLIR_PREFER_NHWC"] = "1"
-        os.environ["ONEFLOW_KERNEL_ENABLE_FUSED_CONV_BIAS"] = "1"
-        os.environ["ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR"] = "1"
-        os.environ["ONEFLOW_KERNEL_CONV_CUTLASS_IMPL_ENABLE_TUNING_WARMUP"] = "1"
-        os.environ["ONEFLOW_KERNEL_CONV_ENABLE_CUTLASS_IMPL"] = "1"
-        os.environ["ONEFLOW_CONV_ALLOW_HALF_PRECISION_ACCUMULATION"] = "1"
-        os.environ["ONEFLOW_MATMUL_ALLOW_HALF_PRECISION_ACCUMULATION"] = "1"
-        os.environ["ONEFLOW_LINEAR_EMBEDDING_SKIP_INIT"] = "1"
+        os.environ.setdefault("ONEFLOW_MLIR_CSE", "1")
+        os.environ.setdefault("ONEFLOW_MLIR_ENABLE_INFERENCE_OPTIMIZATION", "1")
+        os.environ.setdefault("ONEFLOW_MLIR_ENABLE_ROUND_TRIP", "1")
+        os.environ.setdefault("ONEFLOW_MLIR_FUSE_FORWARD_OPS", "1")
+        os.environ.setdefault("ONEFLOW_MLIR_FUSE_OPS_WITH_BACKWARD_IMPL", "1")
+        os.environ.setdefault("ONEFLOW_MLIR_GROUP_MATMUL", "1")
+        os.environ.setdefault("ONEFLOW_MLIR_PREFER_NHWC", "1")
+        os.environ.setdefault("ONEFLOW_KERNEL_ENABLE_FUSED_CONV_BIAS", "1")
+        os.environ.setdefault("ONEFLOW_KERNEL_ENABLE_FUSED_LINEAR", "1")
+        os.environ.setdefault("ONEFLOW_KERNEL_CONV_CUTLASS_IMPL_ENABLE_TUNING_WARMUP", "1")
+        os.environ.setdefault("ONEFLOW_KERNEL_CONV_ENABLE_CUTLASS_IMPL", "1")
+        os.environ.setdefault("ONEFLOW_CONV_ALLOW_HALF_PRECISION_ACCUMULATION", "1")
+        os.environ.setdefault("ONEFLOW_MATMUL_ALLOW_HALF_PRECISION_ACCUMULATION", "1")
+        os.environ.setdefault("ONEFLOW_LINEAR_EMBEDDING_SKIP_INIT", "1")
 
         class OfGraph(flow.nn.Graph):
             def __init__(self):