Set FastMathFlags by default

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/backend/compiler.py

@@ -261,6 +261,7 @@ def make_llir(src, metadata, options):
         if os.getenv("TRITON_INTEL_REDUCE_TRANSPOSE", "0") != "1":
             intel.passes.ttgpuir.add_allocate_shared_memory(pm)
         intel.passes.ttgpuir.add_to_llvmir(pm)
+        intel.set_fast_math(mod)
         passes.convert.add_arith_to_llvmir(pm)
         passes.common.add_canonicalizer(pm)
         passes.common.add_cse(pm)

third_party/intel/lib/TritonIntelGPUTransforms/ScheduleLoad.cpp

@@ -90,15 +90,6 @@ class ScheduleLoadPass
         op->moveAfter(def);
       }
     });
-
-    // HoHo, add fastmath for all
-    // may do this after llvm ir according to user fmath flag
-    mod.walk([&](Operation *op) {
-      if (auto fmIf = dyn_cast<arith::ArithFastMathInterface>(op))
-        op->setAttr(
-            fmIf.getFastMathAttrName(),
-            arith::FastMathFlagsAttr::get(ctx, arith::FastMathFlags::fast));
-    });
   }
 
 private:

third_party/intel/triton_xpu.cc

@@ -215,6 +215,17 @@ void init_triton_intel(py::module &&m) {
     context.loadAllAvailableDialects();
   });
 
+  m.def("set_fast_math", [](mlir::ModuleOp mod) {
+    using namespace mlir;
+    MLIRContext *ctx = mod.getContext();
+    mod.walk([&](Operation *op) {
+      if (auto fmIf = dyn_cast<arith::ArithFastMathInterface>(op))
+        op->setAttr(
+            fmIf.getFastMathAttrName(),
+            arith::FastMathFlagsAttr::get(ctx, arith::FastMathFlags::fast));
+    });
+  });
+
   m.def("set_spv_target_triple", [](llvm::Module *mod) {
     std::string triple = "spir64-unknown-unknown";
     std::string layout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:"