style: add dsl config

Author: AKatydid

graph_net/config/config_agent_backend.yaml

@@ -5,6 +5,9 @@ max_tokens: 16384
 parallel_query_nums: 1
 iterative_query_nums: 2
 
+# support "CUDA", "Triton", future will support "TileLang", etc.
+DSL: "CUDA"
+
 # responses and logs will be saved in <top_save_dir>/<llm_name>
 top_save_dir: "./tmp/llm_cache" 
 

graph_net/torch/backend/agent_backend.py

@@ -19,7 +19,6 @@ def compile(self, *args, **kwargs):
             self.module,
             model_inputs=dummy_input,
             task_name=f"default_task_{self.count_compile}",
-            language="cuda",
         )
         self.count_compile += 1
         return optimized_module

graph_net/torch/backend/agent_ncu/agent_compile.py

@@ -35,21 +35,19 @@ def generate(prompt, system_prompt, llm_query_config: LLMQueryConfig):
         raise RuntimeError(f"LLM query failed with exception: {e}")
 
 
-def optimize(
-    module, model_inputs=None, language: str = "cuda", task_name: str = "default_task"
-):
+def optimize(module, model_inputs=None, task_name: str = "default_task"):
     """Optimize the given PyTorch module using custom DSL operators."""
 
     llm_config = get_llm_config()
     llm_query_config = LLMQueryConfig(**llm_config)
     traced_module = torch.fx.symbolic_trace(module)
 
-    if "cuda" == language:
+    if "cuda" == llm_config.dsl.lower():
         return cuda_optimize(traced_module, model_inputs, task_name, llm_query_config)
-    elif "triton" == language:
+    elif "triton" == llm_config.dsl.lower():
         return torch.compile(module)  # TODO add custom triton optimize
     else:
-        raise NotImplementedError(f"Unsupported language: {language}")
+        raise NotImplementedError(f"Unsupported language: {llm_config.dsl}")
 
     # return the best of optimized models
 

graph_net/torch/backend/agent_utils/query_llm_utils.py

@@ -17,6 +17,8 @@ class LLMQueryConfig:
     parallel_query_nums: int = 1
     iterative_query_nums: int = 10
 
+    dsl: str = "CUDA"
+
     # cache settings
     # responses will be saved in <tmp_llm_cache>/<llm_name>/
     top_save_dir: str = "./tmp_llm_cache"