From 1f8fd525b5faf240a69162422b5697962c1eabce Mon Sep 17 00:00:00 2001
From: Philippe Tillet <ptillet@g.harvard.edu>
Date: Thu, 5 Sep 2019 20:28:00 -0400
Subject: [PATCH] [python] fixed warnings for pybind11 and pytorch

---
 lib/runtime/function.cc          |  2 --
 python/src/bindings.cc           |  2 +-
 python/src/pybind11/functional.h | 24 +++++++++++++++++++-----
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/lib/runtime/function.cc b/lib/runtime/function.cc
index 5454b327f..7c2f42c3c 100644
--- a/lib/runtime/function.cc
+++ b/lib/runtime/function.cc
@@ -220,8 +220,6 @@ std::unique_ptr<driver::module> function::make_bin(ir::module &module, driver::c
   dce.run(module);
   vectorize.run(module);
   dce.run(module);
-
-
   // generate llvm code
   llvm::LLVMContext ctx;
   std::unique_ptr<llvm::Module> llvm(new llvm::Module(module.get_name(), ctx));
diff --git a/python/src/bindings.cc b/python/src/bindings.cc
index 4ef860347..a09a0a7cb 100644
--- a/python/src/bindings.cc
+++ b/python/src/bindings.cc
@@ -493,7 +493,7 @@ extern std::map<size_t, std::shared_ptr<rt::function>> id_fn_map;
 
   oss << std::endl;
   oss << std::endl;
-  oss << "static auto registry = torch::jit::RegisterOperators(\"triton::" << name << "\", &" << name << ");" << std::endl;
+  oss << "static auto registry = torch::RegisterOperators(\"triton::" << name << "\", &" << name << ");" << std::endl;
 
   return {oss.str(), name};
 }
diff --git a/python/src/pybind11/functional.h b/python/src/pybind11/functional.h
index 7a0988ab0..00457e965 100644
--- a/python/src/pybind11/functional.h
+++ b/python/src/pybind11/functional.h
@@ -65,12 +65,26 @@ public:
             }
         };
 
-        value = [hfunc = func_handle(std::move(func))](Args... args) -> Return {
-            gil_scoped_acquire acq;
-            object retval(hfunc.f(std::forward<Args>(args)...));
-            /* Visual studio 2015 parser issue: need parentheses around this expression */
-            return (retval.template cast<Return>());
+        // value = [hfunc = func_handle(std::move(func))](Args... args) -> Return {
+        //     gil_scoped_acquire acq;
+        //     object retval(hfunc.f(std::forward<Args>(args)...));
+        //     /* Visual studio 2015 parser issue: need parentheses around this expression */
+        //     return (retval.template cast<Return>());
+        // };
+
+        struct func_wrapper {
+            func_handle hfunc;
+            func_wrapper(func_handle&& hf): hfunc(std::move(hf)) {}
+            Return operator()(Args... args) const {
+                gil_scoped_acquire acq;
+                object retval(hfunc.f(std::forward<Args>(args)...));
+                /* Visual studio 2015 parser issue: need parentheses around this expression */
+                return (retval.template cast<Return>());
+            }
         };
+
+        value = func_wrapper(func_handle(std::move(func)));
+
         return true;
     }