From a3df6068b4cc17cb11b9143178776b04fa35cd24 Mon Sep 17 00:00:00 2001
From: darkbuck <michael.hliao@gmail.com>
Date: Tue, 15 Aug 2023 21:18:06 -0400
Subject: [PATCH] [BACKEND] Minor fixes found when building triton with LLVM
 17/main branches (#2089)

- These minor fixes are not specific to interface changes from LLVM main
or official llvm-17 branch and can be applied on triton main branch.
- https://github.com/darkbuck/triton/tree/darkbuck/main/llvm-main-branch
has extra changes to build again LLVM main branch build to enable me to
work on other backends on the main branch only. That's the hobby effort
and just FYR.
---
 include/triton/Tools/Sys/GetEnv.hpp                        | 2 ++
 lib/Dialect/TritonGPU/IR/Dialect.cpp                       | 7 ++++---
 .../TritonNvidiaGPU/Transforms/WSMaterialization.cpp       | 2 ++
 python/setup.py                                            | 1 +
 4 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/include/triton/Tools/Sys/GetEnv.hpp b/include/triton/Tools/Sys/GetEnv.hpp
index 5f5040195..53e421ef2 100644
--- a/include/triton/Tools/Sys/GetEnv.hpp
+++ b/include/triton/Tools/Sys/GetEnv.hpp
@@ -24,7 +24,9 @@
 
 #include <algorithm>
 #include <cstdlib>
+#include <set>
 #include <string>
+
 namespace triton {
 
 const std::set<std::string> ENV_VARS = {
diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
index b4daaee79..60da73a22 100644
--- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
+++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -1460,9 +1460,10 @@ struct TritonGPUInferLayoutInterface
     return success();
   }
 
-  LogicalResult inferDotOpEncoding(Attribute operandEncoding, unsigned opIdx,
-                                   Attribute retEncoding,
-                                   Optional<Location> location) const override {
+  LogicalResult
+  inferDotOpEncoding(Attribute operandEncoding, unsigned opIdx,
+                     Attribute retEncoding,
+                     std::optional<Location> location) const override {
     auto mmaRetEncoding = retEncoding.dyn_cast<MmaEncodingAttr>();
     if (mmaRetEncoding && mmaRetEncoding.isHopper()) {
       // TODO: support gmma when A/B does not reside in shared memory
diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/WSMaterialization.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/WSMaterialization.cpp
index bd92ba372..9ebc78497 100644
--- a/lib/Dialect/TritonNvidiaGPU/Transforms/WSMaterialization.cpp
+++ b/lib/Dialect/TritonNvidiaGPU/Transforms/WSMaterialization.cpp
@@ -30,6 +30,8 @@
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Utility.h"
 
+#include <set>
+
 using namespace mlir;
 namespace ttg = triton::gpu;
 namespace ttng = triton::nvidia_gpu;
diff --git a/python/setup.py b/python/setup.py
index 19df06997..18764ec13 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -335,6 +335,7 @@ setup(
             "matplotlib",
             "pandas",
             "tabulate",
+            "torch",
         ],
     },
 )