#1085: Allow the user to apply LLVM optimizations (-O3, -Oz) when lif…

…ting to LLVM

.build_number

@@ -1 +1 @@
-1519
+1520

src/libtriton/api/api.cpp

@@ -1373,17 +1373,17 @@ namespace triton {
 
   /* Lifters engine API ================================================================================= */
 
-  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname) {
+  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname, bool optimize) {
     this->checkLifting();
     #ifdef TRITON_LLVM_INTERFACE
-    return this->lifting->liftToLLVM(stream, node, fname);
+    return this->lifting->liftToLLVM(stream, node, fname, optimize);
     #endif
     throw triton::exceptions::API("API::liftToLLVM(): Triton not built with LLVM");
   }
 
 
-  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname) {
-    return this->liftToLLVM(stream, expr->getAst(), fname);
+  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname, bool optimize) {
+    return this->liftToLLVM(stream, expr->getAst(), fname, optimize);
   }
 
 

src/libtriton/ast/llvm/tritonToLLVM.cpp

@@ -81,7 +81,7 @@ namespace triton {
     }
 
 
-    std::shared_ptr<llvm::Module> TritonToLLVM::convert(const triton::ast::SharedAbstractNode& node, const char* fname) {
+    std::shared_ptr<llvm::Module> TritonToLLVM::convert(const triton::ast::SharedAbstractNode& node, const char* fname, bool optimize) {
       std::unordered_map<triton::ast::SharedAbstractNode, llvm::Value*> results;
 
       /* Create the LLVM function */
@@ -98,6 +98,16 @@ namespace triton {
       /* Create the return instruction */
       this->llvmIR.CreateRet(results.at(node));
 
+      /* Apply LLVM optimizations (-03 -Oz) if enabled */
+      if (optimize) {
+        llvm::legacy::PassManager pm;
+        llvm::PassManagerBuilder pmb;
+        pmb.OptLevel = 3;
+        pmb.SizeLevel = 2;
+        pmb.populateModulePassManager(pm);
+        pm.run(*this->llvmModule);
+      }
+
       return this->llvmModule;
     }
 

src/libtriton/bindings/python/objects/pyTritonContext.cpp

@@ -286,11 +286,11 @@ Returns true if the taint engine is enabled.
 - <b>bool isThumb(void)</b><br>
 Returns true if execution mode is Thumb (only valid for ARM32).
 
-- <b>string liftToLLVM(\ref py_AstNode_page node, string fname="__triton")</b><br>
-Lifts an AST node and all its references to LLVM IR. `fname` is the name of the LLVM function, by default it's `__triton`.
+- <b>string liftToLLVM(\ref py_AstNode_page node, string fname="__triton", bool optimize=False)</b><br>
+Lifts an AST node and all its references to LLVM IR. `fname` is the name of the LLVM function, by default it's `__triton`. If `optimize` is true, perform optimizations (-O3 -Oz).
 
-- <b>string liftToLLVM(\ref py_SymbolicExpression_page expr)</b><br>
-Lifts a symbolic expression and all its references to LLVM IR.
+- <b>string liftToLLVM(\ref py_SymbolicExpression_page expr, string fname="__triton", bool optimize=False)</b><br>
+Lifts a symbolic expression and all its references to LLVM IR. `fname` is the name of the LLVM function, by default it's `__triton`. If `optimize` is true, perform optimizations (-O3 -Oz).
 
 - <b>string liftToPython(\ref py_SymbolicExpression_page expr)</b><br>
 Lifts a symbolic expression and all its references to Python format.
@@ -2232,17 +2232,19 @@ namespace triton {
 
 
       static PyObject* TritonContext_liftToLLVM(PyObject* self, PyObject* args, PyObject* kwargs) {
-        PyObject* node  = nullptr;
-        PyObject* fname = nullptr;
+        PyObject* node      = nullptr;
+        PyObject* fname     = nullptr;
+        PyObject* optimize  = nullptr;
 
         static char* keywords[] = {
           (char*)"node",
           (char*)"fname",
+          (char*)"optimize",
           nullptr
         };
 
         /* Extract keywords */
-        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OO", keywords, &node, &fname) == false) {
+        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OOO", keywords, &node, &fname, &optimize) == false) {
           return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Invalid number of arguments");
         }
 
@@ -2252,16 +2254,22 @@ namespace triton {
         if (fname != nullptr && !PyStr_Check(fname))
           return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a string as fname parameter.");
 
+        if (optimize != nullptr && !PyBool_Check(optimize))
+          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a boolean as optimize parameter.");
+
         if (fname == nullptr)
           fname = PyStr_FromString("__triton");
 
+        if (optimize == nullptr)
+          optimize = PyLong_FromUint32(false);
+
         try {
           std::ostringstream stream;
           if (PySymbolicExpression_Check(node)) {
-            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PySymbolicExpression_AsSymbolicExpression(node), PyStr_AsString(fname));
+            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PySymbolicExpression_AsSymbolicExpression(node), PyStr_AsString(fname), PyLong_AsBool(optimize));
           }
           else {
-            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PyAstNode_AsAstNode(node), PyStr_AsString(fname));
+            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PyAstNode_AsAstNode(node), PyStr_AsString(fname), PyLong_AsBool(optimize));
           }
           return xPyString_FromString(stream.str().c_str());
         }

src/libtriton/engines/lifters/liftingToLLVM.cpp

@@ -22,21 +22,21 @@ namespace triton {
       }
 
 
-      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname) {
-        this->liftToLLVM(stream, expr->getAst(), fname);
+      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname, bool optimize) {
+        this->liftToLLVM(stream, expr->getAst(), fname, optimize);
         return stream;
       }
 
 
-      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname) {
+      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname, bool optimize) {
         /* The LLVM context */
         llvm::LLVMContext context;
 
         /* The lifter Triton -> LLVM */
         triton::ast::TritonToLLVM lifter(context);
 
         /* Lift AST to LLVM IR */
-        auto llvmModule = lifter.convert(node, fname);
+        auto llvmModule = lifter.convert(node, fname, optimize);
 
         /* Print the LLVM module into the stream */
         std::string dump;

src/libtriton/includes/triton/api.hpp

@@ -673,10 +673,10 @@ namespace triton {
         /* Lifters engine API ================================================================================= */
 
         //! [**lifting api**] - Lifts an AST and all its references to LLVM format. `fname` represents the name of the LLVM function.
-        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname="__triton");
+        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname="__triton", bool optimize=false);
 
         //! [**lifting api**] - Lifts a symbolic expression and all its references to LLVM format. `fname` represents the name of the LLVM function.
-        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname="__triton");
+        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname="__triton", bool optimize=false);
 
         //! [**lifting api**] - Lifts a symbolic expression and all its references to Python format.
         TRITON_EXPORT std::ostream& liftToPython(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr);

src/libtriton/includes/triton/liftingToLLVM.hpp

@@ -53,10 +53,10 @@ namespace triton {
           TRITON_EXPORT LiftingToLLVM();
 
           //! Lifts a symbolic expression and all its references to LLVM format. `fname` represents the name of the LLVM function.
-          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname="__triton");
+          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname="__triton", bool optimize=false);
 
           //! Lifts a abstract node and all its references to LLVM format. `fname` represents the name of the LLVM function.
-          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname="__triton");
+          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname="__triton", bool optimize=false);
       };
 
     /*! @} End of lifters namespace */

src/libtriton/includes/triton/tritonToLLVM.hpp

@@ -19,7 +19,9 @@
 
 #include <llvm/IR/IRBuilder.h>
 #include <llvm/IR/LLVMContext.h>
+#include <llvm/IR/LegacyPassManager.h>
 #include <llvm/IR/Module.h>
+#include <llvm/Transforms/IPO/PassManagerBuilder.h>
 
 
 
@@ -65,7 +67,7 @@ namespace triton {
         TRITON_EXPORT TritonToLLVM(llvm::LLVMContext& llvmContext);
 
         //! Lifts a symbolic expression and all its references to LLVM format. `fname` represents the name of the LLVM function.
-        TRITON_EXPORT std::shared_ptr<llvm::Module> convert(const triton::ast::SharedAbstractNode& node, const char* fname="__triton");
+        TRITON_EXPORT std::shared_ptr<llvm::Module> convert(const triton::ast::SharedAbstractNode& node, const char* fname="__triton", bool optimize=false);
     };
 
   /*! @} End of ast namespace */

src/testers/unittests/test_ast_representation.py

@@ -269,4 +269,4 @@ def test_lifting(self):
             ]
 
             for n in nodes:
-                self.assertNotEqual(len(self.ctx.liftToLLVM(n, fname="test")), 0)
+                self.assertNotEqual(len(self.ctx.liftToLLVM(n, fname="test", optimize=True)), 0)