From e1e9d731872fe928c7f32015d575401c667145b0 Mon Sep 17 00:00:00 2001
From: Digant Desai <digantdesai@meta.com>
Date: Thu, 23 Jan 2025 22:51:25 -0600
Subject: [PATCH] Update default target selection logic

First try to lookup the Target in the given module.
If it doesn't work, use the default target. And set it in the module.

Rationale: Issue #207
---
 python/src/llvm.cc | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/src/llvm.cc b/python/src/llvm.cc
index e4be9846bcc4..e6a7d51e86e9 100644
--- a/python/src/llvm.cc
+++ b/python/src/llvm.cc
@@ -57,13 +57,17 @@ std::unique_ptr<TargetMachine>
 createTargetMachine(llvm::Module *module, std::string proc,
                     bool enable_fp_fusion, const std::string &features,
                     bool enable_fast_math = false) {
-  auto triple = getDefaultTargerOrProcessTriple();
-  module->setTargetTriple(triple);
   std::string error;
   auto target =
       llvm::TargetRegistry::lookupTarget(module->getTargetTriple(), error);
   if (!target) {
-    throw std::runtime_error("target lookup error: " + error);
+    // Try to get the default target triple.
+    auto triple = getDefaultTargerOrProcessTriple();
+    target = llvm::TargetRegistry::lookupTarget(triple, error);
+    if (!target) {
+      throw std::runtime_error("target lookup error: " + error);
+    }
+    module->setTargetTriple(triple);
   }
   llvm::TargetOptions opt;
   bool disableLLVMOpt = mlir::triton::tools::getBoolEnv("DISABLE_LLVM_OPT");