add env base

PaddlePaddle · Jan 17, 2024 · 9b72523 · 9b72523
1 parent d9737c6
commit 9b72523
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/padiff/cinn_diff/run.py b/padiff/cinn_diff/run.py
@@ -26,5 +26,24 @@ def run(run_script, base_env, cinn_env):
 
 
 if __name__ == "__main__":
+    _base_env = {
+            "CUDA_VISIBLE_DEVICES": "7",
+            "NVIDIA_TF32_OVERRIDE": "1",
+            "CUDA_LAUNCH_BLOCKING": "1",
+            "FLAGS_save_static_runtime_data": "1",
+            "FLAGS_static_runtime_data_save_path": "./",
+            "FLAGS_cudnn_deterministc": "1",
+            "FLAGS_cinn_cudnn_deterministc": "1",
+            "FLAGS_prim_all": "true",
+        }
+    _cinn_env = {
+            "FLAGS_use_cinn": "1",
+            "FLAGS_deny_cinn_ops": "",
+            "FLAGS_use_reduce_split_pass": "1",
+            "FLAGS_nvrtc_compile_to_cubin": "0",
+            "FLAGS_cinn_use_op_fusion": "1",
+            "FLAGS_cinn_parallel_compile_size": "8",
+            "FLAGS_cinn_pass_visualize_dir": "",
+        }
     run_script = "/root/dev/PaddleNLP/model_zoo/bert/run_bert.sh"
-    run(run_script, None, None)
+    run(run_script, _base_env, _cinn_env)