|
10 | 10 | #include <imex/Dialect/PTensor/IR/PTensorOps.h> |
11 | 11 | #include <mlir/Dialect/LLVMIR/LLVMDialect.h> |
12 | 12 |
|
| 13 | +#include <iostream> |
| 14 | + |
13 | 15 | static tbb::concurrent_bounded_queue<Runable::ptr_type> _deferred; |
14 | 16 |
|
15 | 17 | void push_runable(Runable::ptr_type && r) |
@@ -63,42 +65,95 @@ void process_promises() |
63 | 65 | ::mlir::OpBuilder builder(&jit._context); |
64 | 66 | auto loc = builder.getUnknownLoc(); |
65 | 67 | jit::IdValueMap ivp; |
66 | | - ::mlir::Value ret_value; |
67 | 68 |
|
68 | 69 | // Create a MLIR module |
69 | 70 | auto module = builder.create<::mlir::ModuleOp>(loc); |
70 | 71 | // Create a func |
71 | 72 | auto dtype = builder.getI64Type(); |
72 | | - llvm::SmallVector<int64_t> shape(1, -1); //::mlir::ShapedType::kDynamicSize); |
73 | | - auto rrtype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get(shape, dtype), true); // llvm::SmallVector<int64_t>() |
74 | | - auto funcType = builder.getFunctionType({}, rrtype); |
| 73 | + // create dummy type, we'll replace it with the actual type later |
| 74 | + auto dummyFuncType = builder.getFunctionType({}, dtype); |
75 | 75 | std::string fname("ddpt_jit"); |
76 | | - auto function = builder.create<::mlir::func::FuncOp>(loc, fname, funcType); |
77 | | - // request generation of c-wrapper function |
78 | | - function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(), ::mlir::UnitAttr::get(&jit._context)); |
| 76 | + auto function = builder.create<::mlir::func::FuncOp>(loc, fname, dummyFuncType); |
79 | 77 | // create function entry block |
80 | 78 | auto &entryBlock = *function.addEntryBlock(); |
81 | 79 | // Set the insertion point in the builder to the beginning of the function body |
82 | 80 | builder.setInsertionPointToStart(&entryBlock); |
83 | | - |
| 81 | + // we need to keep runables/deferred/futures alive until we set their values below |
| 82 | + std::vector<Runable::ptr_type> runables; |
| 83 | + |
84 | 84 | while(true) { |
85 | 85 | Runable::ptr_type d; |
86 | 86 | _deferred.pop(d); |
87 | 87 | if(d) { |
88 | | - d->run(); |
89 | | - ret_value = d->generate_mlir(builder, loc, ivp); |
90 | | - d.reset(); |
| 88 | + // d->run(); |
| 89 | + (void) d->generate_mlir(builder, loc, ivp); |
| 90 | + // keep alive for later set_value |
| 91 | + runables.push_back(std::move(d)); |
| 92 | + //d.reset(); |
91 | 93 | } else { |
92 | 94 | break; |
93 | 95 | } |
94 | 96 | } |
95 | 97 |
|
96 | | - (void)builder.create<::mlir::func::ReturnOp>(loc, ret_value); |
| 98 | + // Now we have to define the return type as a ValueRange of all arrays which we have created |
| 99 | + // (runnables have put them into ivp) |
| 100 | + // We also compute the total size of the struct llvm created for this return type |
| 101 | + // llvm will basically return a struct with all the arrays as members, each of type JIT::MemRefDescriptor |
| 102 | + |
| 103 | + // Need a container to put all return values, will be used to construct TypeRange |
| 104 | + std::vector<::mlir::Type> ret_types; |
| 105 | + ret_types.reserve(ivp.size()); |
| 106 | + std::vector<::mlir::Value> ret_values; |
| 107 | + ret_types.reserve(ivp.size()); |
| 108 | + std::unordered_map<id_type, uint64_t> rank_map; |
| 109 | + // here we store the total size of the llvm struct |
| 110 | + uint64_t sz = 0; |
| 111 | + for(auto & v : ivp) { |
| 112 | + auto value = v.second.first; |
| 113 | + // append the type and array/value |
| 114 | + ret_types.push_back(value.getType()); |
| 115 | + ret_values.push_back(value); |
| 116 | + auto ptt = value.getType().dyn_cast<::imex::ptensor::PTensorType>(); |
| 117 | + assert(ptt); |
| 118 | + auto rank = ptt.getRtensor().getShape().size(); |
| 119 | + rank_map[v.first] = rank; |
| 120 | + // add sizeof(MemRefDescriptor<elementtype, rank>) to sz |
| 121 | + sz += 3 + 2 * rank; |
| 122 | + } |
| 123 | + ::mlir::TypeRange ret_tr(ret_types); |
| 124 | + ::mlir::ValueRange ret_vr(ret_values); |
| 125 | + |
| 126 | + // add return statement |
| 127 | + auto ret_value = builder.create<::mlir::func::ReturnOp>(loc, ret_vr); |
| 128 | + // Define and assign correct function type |
| 129 | + auto funcTypeAttr = ::mlir::TypeAttr::get(builder.getFunctionType({}, ret_tr)); |
| 130 | + function.setFunctionTypeAttr(funcTypeAttr); |
| 131 | + // also request generation of c-wrapper function |
| 132 | + function->setAttr(::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName(), ::mlir::UnitAttr::get(&jit._context)); |
97 | 133 | // add the function to the module |
98 | 134 | module.push_back(function); |
99 | 135 | module.dump(); |
100 | 136 | // finally compile and run the module |
101 | | - if(jit.run(module, fname)) throw std::runtime_error("failed running jit"); |
| 137 | + assert(sizeof(intptr_t) == sizeof(void*)); |
| 138 | + intptr_t * output = new intptr_t[sz]; |
| 139 | + std::cout << ivp.size() << " sz: " << sz << std::endl; |
| 140 | + if(jit.run(module, fname, output)) throw std::runtime_error("failed running jit"); |
| 141 | + |
| 142 | + size_t pos = 0; |
| 143 | + for(auto & v : ivp) { |
| 144 | + auto value = v.second.first; |
| 145 | + auto rank = rank_map[v.first]; |
| 146 | + void * allocated = (void*)output[pos]; |
| 147 | + void * aligned = (void*)output[pos+1]; |
| 148 | + intptr_t offset = output[pos+2]; |
| 149 | + intptr_t * sizes = output + pos + 3; |
| 150 | + intptr_t * stride = output + pos + 3 + rank; |
| 151 | + pos += 3 + 2 * rank; |
| 152 | + v.second.second(rank, allocated, aligned, offset, sizes, stride); |
| 153 | + } |
| 154 | + |
| 155 | + // finally release all our runables/tasks/deferred/futures |
| 156 | + runables.clear(); |
102 | 157 | } |
103 | 158 |
|
104 | 159 | void sync() |
|
0 commit comments