@@ -81,99 +81,107 @@ class DepManager
8181
8282void process_promises ()
8383{
84+ bool done = false ;
8485 jit::JIT jit;
85- ::mlir::OpBuilder builder (&jit._context );
86- auto loc = builder.getUnknownLoc ();
87- jit::IdValueMap ivp;
88-
89- // Create a MLIR module
90- auto module = builder.create <::mlir::ModuleOp>(loc);
91- // Create a func
92- auto dtype = builder.getI64Type ();
93- // create dummy type, we'll replace it with the actual type later
94- auto dummyFuncType = builder.getFunctionType ({}, dtype);
95- std::string fname (" ddpt_jit"
96- auto function = builder.create <::mlir::func::FuncOp>(loc, fname, dummyFuncType);
97- // create function entry block
98- auto &entryBlock = *function.addEntryBlock ();
99- // Set the insertion point in the builder to the beginning of the function body
100- builder.setInsertionPointToStart (&entryBlock);
101- // we need to keep runables/deferred/futures alive until we set their values below
102- std::vector<Runable::ptr_type> runables;
103-
104- while (true ) {
105- Runable::ptr_type d;
106- _deferred.pop (d);
107- if (d) {
108- // d->run();
109- (void ) d->generate_mlir (builder, loc, ivp);
110- // keep alive for later set_value
111- runables.push_back (std::move (d));
112- // d.reset();
113- } else {
114- break ;
115- }
116- }
11786
118- // Now we have to define the return type as a ValueRange of all arrays which we have created
119- // (runnables have put them into ivp)
120- // We also compute the total size of the struct llvm created for this return type
121- // llvm will basically return a struct with all the arrays as members, each of type JIT::MemRefDescriptor
122-
123- // Need a container to put all return values, will be used to construct TypeRange
124- std::vector<::mlir::Type> ret_types;
125- ret_types.reserve (ivp.size ());
126- std::vector<::mlir::Value> ret_values;
127- ret_types.reserve (ivp.size ());
128- std::unordered_map<id_type, uint64_t > rank_map;
129- // here we store the total size of the llvm struct
130- uint64_t sz = 0 ;
131- for (auto & v : ivp) {
132- auto value = v.second .first ;
133- // append the type and array/value
134- ret_types.push_back (value.getType ());
135- ret_values.push_back (value);
136- auto ptt = value.getType ().dyn_cast <::imex::ptensor::PTensorType>();
137- assert (ptt);
138- auto rank = ptt.getRtensor ().getShape ().size ();
139- rank_map[v.first ] = rank;
140- // add sizeof(MemRefDescriptor<elementtype, rank>) to sz
141- sz += 3 + 2 * rank;
142- }
143- ::mlir::TypeRange ret_tr (ret_types);
144- ::mlir::ValueRange ret_vr (ret_values);
145-
146- // add return statement
147- auto ret_value = builder.create <::mlir::func::ReturnOp>(loc, ret_vr);
148- // Define and assign correct function type
149- auto funcTypeAttr = ::mlir::TypeAttr::get (builder.getFunctionType ({}, ret_tr));
150- function.setFunctionTypeAttr (funcTypeAttr);
151- // also request generation of c-wrapper function
152- function->setAttr (::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName (), ::mlir::UnitAttr::get (&jit._context ));
153- // add the function to the module
154- module .push_back (function);
155- module .dump ();
156- // finally compile and run the module
157- assert (sizeof (intptr_t ) == sizeof (void *));
158- intptr_t * output = new intptr_t [sz];
159- std::cout << ivp.size () << " sz: "
160- if (jit.run (module , fname, output)) throw std::runtime_error (" failed running jit"
87+ do {
88+ ::mlir::OpBuilder builder (&jit._context );
89+ auto loc = builder.getUnknownLoc ();
90+ jit::IdValueMap ivp;
16191
162- size_t pos = 0 ;
163- for (auto & v : ivp) {
164- auto value = v.second .first ;
165- auto rank = rank_map[v.first ];
166- void * allocated = (void *)output[pos];
167- void * aligned = (void *)output[pos+1 ];
168- intptr_t offset = output[pos+2 ];
169- intptr_t * sizes = output + pos + 3 ;
170- intptr_t * stride = output + pos + 3 + rank;
171- pos += 3 + 2 * rank;
172- v.second .second (rank, allocated, aligned, offset, sizes, stride);
173- }
92+ // Create a MLIR module
93+ auto module = builder.create <::mlir::ModuleOp>(loc);
94+ // Create a func
95+ auto dtype = builder.getI64Type ();
96+ // create dummy type, we'll replace it with the actual type later
97+ auto dummyFuncType = builder.getFunctionType ({}, dtype);
98+ std::string fname (" ddpt_jit"
99+ auto function = builder.create <::mlir::func::FuncOp>(loc, fname, dummyFuncType);
100+ // create function entry block
101+ auto &entryBlock = *function.addEntryBlock ();
102+ // Set the insertion point in the builder to the beginning of the function body
103+ builder.setInsertionPointToStart (&entryBlock);
104+ // we need to keep runables/deferred/futures alive until we set their values below
105+ std::vector<Runable::ptr_type> runables;
106+
107+ while (true ) {
108+ Runable::ptr_type d;
109+ _deferred.pop (d);
110+ if (d) {
111+ if (d->generate_mlir (builder, loc, ivp)) {
112+ d.reset ();
113+ break ;
114+ };
115+ // keep alive for later set_value
116+ runables.push_back (std::move (d));
117+ } else {
118+ // signals system shutdown
119+ done = true ;
120+ break ;
121+ }
122+ }
174123
175- // finally release all our runables/tasks/deferred/futures
176- runables.clear ();
124+ if (runables.empty ()) continue ;
125+
126+ // Now we have to define the return type as a ValueRange of all arrays which we have created
127+ // (runnables have put them into ivp when generating mlir)
128+ // We also compute the total size of the struct llvm created for this return type
129+ // llvm will basically return a struct with all the arrays as members, each of type JIT::MemRefDescriptor
130+
131+ // Need a container to put all return values, will be used to construct TypeRange
132+ std::vector<::mlir::Type> ret_types;
133+ ret_types.reserve (ivp.size ());
134+ std::vector<::mlir::Value> ret_values;
135+ ret_types.reserve (ivp.size ());
136+ std::unordered_map<id_type, uint64_t > rank_map;
137+ // here we store the total size of the llvm struct
138+ uint64_t sz = 0 ;
139+ for (auto & v : ivp) {
140+ auto value = v.second .first ;
141+ // append the type and array/value
142+ ret_types.push_back (value.getType ());
143+ ret_values.push_back (value);
144+ auto ptt = value.getType ().dyn_cast <::imex::ptensor::PTensorType>();
145+ assert (ptt);
146+ auto rank = ptt.getRtensor ().getShape ().size ();
147+ rank_map[v.first ] = rank;
148+ // add sizeof(MemRefDescriptor<elementtype, rank>) to sz
149+ sz += 3 + 2 * rank;
150+ }
151+ ::mlir::TypeRange ret_tr (ret_types);
152+ ::mlir::ValueRange ret_vr (ret_values);
153+
154+ // add return statement
155+ auto ret_value = builder.create <::mlir::func::ReturnOp>(loc, ret_vr);
156+ // Define and assign correct function type
157+ auto funcTypeAttr = ::mlir::TypeAttr::get (builder.getFunctionType ({}, ret_tr));
158+ function.setFunctionTypeAttr (funcTypeAttr);
159+ // also request generation of c-wrapper function
160+ function->setAttr (::mlir::LLVM::LLVMDialect::getEmitCWrapperAttrName (), ::mlir::UnitAttr::get (&jit._context ));
161+ // add the function to the module
162+ module .push_back (function);
163+ module .dump ();
164+
165+ // compile and run the module
166+ assert (sizeof (intptr_t ) == sizeof (void *));
167+ intptr_t * output = new intptr_t [sz];
168+ std::cout << ivp.size () << " sz: "
169+ if (jit.run (module , fname, output)) throw std::runtime_error (" failed running jit"
170+
171+ // push results to fulfill promises
172+ size_t pos = 0 ;
173+ for (auto & v : ivp) {
174+ auto value = v.second .first ;
175+ auto rank = rank_map[v.first ];
176+ void * allocated = (void *)output[pos];
177+ void * aligned = (void *)output[pos+1 ];
178+ intptr_t offset = output[pos+2 ];
179+ intptr_t * sizes = output + pos + 3 ;
180+ intptr_t * stride = output + pos + 3 + rank;
181+ pos += 3 + 2 * rank;
182+ v.second .second (rank, allocated, aligned, offset, sizes, stride);
183+ }
184+ } while (!done);
177185}
178186
179187void sync ()
0 commit comments