initial link deferredrunner->jit

Author: fschlimb

setup.py

@@ -33,13 +33,14 @@ def build_cmake(self, ext):
         cmake_args = [
             '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + str(extdir.parent.absolute()),
             '-DCMAKE_BUILD_TYPE=' + config,
-            '-DCMAKE_VERBOSE_MAKEFILE=ON'
+            '-DCMAKE_VERBOSE_MAKEFILE=ON',
+            '-G=Ninja',
         ]
 
         # example of build args
         build_args = [
             '--config', config,
-            '--', '-j8'
+            # '--', '-j8'
         ]
 
         os.chdir(str(build_temp))

src/Creator.cpp

@@ -4,6 +4,9 @@
 #include "ddptensor/Deferred.hpp"
 #include "ddptensor/Factory.hpp"
 
+#include <imex/Dialect/PTensor/IR/PTensorOps.h>
+#include <mlir/IR/Builders.h>
+
 namespace x {
 
     template<typename T>
@@ -141,10 +144,25 @@ struct DeferredArange : public Deferred
         : _start(start), _end(end), _step(step), _dtype(dtype)
     {}
 
-    void run()
+    void run() override
     {
         set_value(std::move(TypeDispatch<x::Creator>(_dtype, _start, _end, _step)));
     };
+    
+    ::mlir::Value generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::IdValueMap & ivm) override
+    {
+        // create start, stop and step
+        auto start = jit::createI64(loc, builder, _start);
+        auto end = jit::createI64(loc, builder, _end);
+        auto step = jit::createI64(loc, builder, _step);
+        // create arange
+        auto dtype = builder.getI64Type();
+        llvm::SmallVector<int64_t> shape(1, -1); //::mlir::ShapedType::kDynamicSize);
+        auto artype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get(shape, dtype), true);
+        auto ar = builder.create<::imex::ptensor::ARangeOp>(loc, artype, start, end, step, true);
+        ivm[_guid] = ar;
+        return ar;
+    }
 
     FactoryId factory() const
     {

src/Deferred.cpp

@@ -1,11 +1,15 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
-#include <oneapi/tbb/concurrent_queue.h>
 #include "include/ddptensor/Deferred.hpp"
 #include "include/ddptensor/Transceiver.hpp"
 #include "include/ddptensor/Mediator.hpp"
 #include "include/ddptensor/Registry.hpp"
 
+#include <oneapi/tbb/concurrent_queue.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <imex/Dialect/PTensor/IR/PTensorOps.h>
+#include <mlir/Dialect/LLVMIR/LLVMDialect.h>
+
 static tbb::concurrent_bounded_queue<Runable::ptr_type> _deferred;
 
 void push_runable(Runable::ptr_type && r)
@@ -55,13 +59,46 @@ void Runable::fini()
 
 void process_promises()
 {
+    jit::JIT jit;
+    ::mlir::OpBuilder builder(&jit._context);
+    auto loc = builder.getUnknownLoc();
+    jit::IdValueMap ivp;
+    ::mlir::Value ret_value;
+
+    // Create a MLIR module
+    auto module = builder.create<::mlir::ModuleOp>(loc);
+    // Create a func
+    auto dtype = builder.getI64Type();
+    llvm::SmallVector<int64_t> shape(1, -1); //::mlir::ShapedType::kDynamicSize);
+    auto rrtype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get(shape, dtype), true); // llvm::SmallVector<int64_t>()
+    auto funcType = builder.getFunctionType({}, rrtype);
+    std::string fname("ddpt_jit");
+    auto function = builder.create<::mlir::func::FuncOp>(loc, fname, funcType);
+    // request generation of c-wrapper function
+    function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(), ::mlir::UnitAttr::get(&jit._context));
+    // create function entry block
+    auto &entryBlock = *function.addEntryBlock();
+    // Set the insertion point in the builder to the beginning of the function body
+    builder.setInsertionPointToStart(&entryBlock);
+    
     while(true) {
         Runable::ptr_type d;
         _deferred.pop(d);
-        if(d) d->run();
-        else break;
-        d.reset();
+        if(d) {
+            d->run();
+            ret_value = d->generate_mlir(builder, loc, ivp);
+            d.reset();
+        } else {
+            break;
+        }
     }
+
+    (void)builder.create<::mlir::func::ReturnOp>(loc, ret_value);
+    // add the function to the module
+    module.push_back(function);
+    module.dump();
+    // finally compile and run the module
+    if(jit.run(module, fname)) throw std::runtime_error("failed running jit");
 }
 
 void sync()

src/include/ddptensor/Deferred.hpp

@@ -4,6 +4,7 @@
 
 #include "tensor_i.hpp"
 #include "Registry.hpp"
+#include "jit/mlir.hpp"
 
 extern void process_promises();
 extern void sync();
@@ -12,7 +13,10 @@ struct Runable
 {
     using ptr_type = std::unique_ptr<Runable>;
     virtual ~Runable() {};
+    /// actually execute, a deferred will set value of future
     virtual void run() = 0;
+    /// generate MLIR code for jit
+    virtual ::mlir::Value generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location, jit::IdValueMap & ivm) = 0;
     virtual FactoryId factory() const = 0;
     virtual void defer(ptr_type &&);
     static void fini();
@@ -29,13 +33,20 @@ struct DeferredT : public P, public Runable
 
     DeferredT() = default;
     DeferredT(const DeferredT<P, F> &) = delete;
+
+    // FIXME: from Runable but should be in most derived classes
+    ::mlir::Value generate_mlir(::mlir::OpBuilder &, ::mlir::Location, jit::IdValueMap &) override
+    {
+        return {};
+    };
 };
 
 struct Deferred : public DeferredT<tensor_i::promise_type, tensor_i::future_type>
 {
     using ptr_type = std::unique_ptr<Deferred>;
     id_type _guid = Registry::NOGUID;
     future_type get_future();
+    // from Runable
     void defer(Runable::ptr_type &&);
 };
 
@@ -98,6 +109,13 @@ struct DeferredLambda : public Runable
         _l();
     }
 
+    // FIXME: from Runable but should be in most derived classes
+    ::mlir::Value generate_mlir(::mlir::OpBuilder &, ::mlir::Location, jit::IdValueMap &) override
+    {
+        throw(std::runtime_error("No MLIR support for DeferredLambda."));
+        return {};
+    };
+
     FactoryId factory() const
     {
         throw(std::runtime_error("No Factory for DeferredLambda."));

src/include/ddptensor/jit/mlir.hpp

@@ -2,17 +2,27 @@
 
 #pragma once
 
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/MLIRContext.h"
-#include "mlir/Pass/PassManager.h"
+#include "ddptensor/UtilsAndTypes.hpp"
+
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/MLIRContext.h>
+#include <mlir/Pass/PassManager.h>
+#include <mlir/IR/Builders.h>
+
+#include <unordered_map>
 
 namespace jit {
 
+using IdValueMap = std::unordered_map<id_type, ::mlir::Value>;
+
 // initialize jit
 void init();
 
 void ttt();
 
+// create a constant integer with given value
+extern ::mlir::Value createI64(const ::mlir::Location & loc, ::mlir::OpBuilder & builder, int64_t val);
+
 // A class to manage the MLIR business (compilation and execution).
 // Just a stub for now, will need to be extended with paramters and maybe more.
 class JIT {

src/jit/mlir.cpp

@@ -57,7 +57,7 @@ static ::mlir::Type makeSignlessType(::mlir::Type type)
     return type;
 }
 
-auto createI64(const ::mlir::Location & loc, ::mlir::OpBuilder & builder, int64_t val)
+::mlir::Value createI64(const ::mlir::Location & loc, ::mlir::OpBuilder & builder, int64_t val)
 {
     auto attr = builder.getI64IntegerAttr(val);
     return builder.create<::mlir::arith::ConstantOp>(loc, attr).getResult();