| 168 | } |
| 169 | |
| 170 | std::vector<Express::VARP> WhileModule::onForward(const std::vector<Express::VARP>& inputsI) { |
| 171 | std::vector<Express::VARP> bodyInputs(mInfo->mBodyInputNumber); |
| 172 | auto& inputs = inputsI; |
| 173 | int step = 0; |
| 174 | std::vector<Express::VARP> outputs(mInfo->mOutputNumber); |
| 175 | if (mCond == nullptr) { |
| 176 | auto limit = inputs[0]->readMap<int>()[0]; |
| 177 | int cond = inputs[1]->readMap<int>()[0]; |
| 178 | // Body Input: 2 + N, Body Output: 1 + N + K, Op output: N + K |
| 179 | int N = mInfo->mBodyInputNumber - 2; |
| 180 | int K = mInfo->mOutputNumber - N; |
| 181 | std::vector<std::vector<VARP>> spans(K); |
| 182 | std::vector<VARP> bodyOutputs; |
| 183 | for (int i=0; i<N; ++i) { |
| 184 | outputs[i] = inputs[i+2]; |
| 185 | } |
| 186 | if (limit > 0 && cond > 0) { |
| 187 | bodyInputs = inputs; |
| 188 | bodyInputs[0] = _Input({}, NCHW, halide_type_of<int>()); |
| 189 | while (step < limit && cond > 0) { |
| 190 | bodyInputs[0]->writeMap<int>()[0] = step; |
| 191 | bodyOutputs = mBody->onForward(bodyInputs); |
| 192 | if (bodyOutputs.empty()) { |
| 193 | // Has Error |
| 194 | return {}; |
| 195 | } |
| 196 | for (int i=0; i<N; ++i) { |
| 197 | bodyInputs[i + 2] = bodyOutputs[i + 1]; |
| 198 | } |
| 199 | for (int i=0; i<K; ++i) { |
| 200 | spans[i].emplace_back(bodyOutputs[1+N+i]); |
| 201 | } |
| 202 | step++; |
| 203 | cond = bodyOutputs[0]->readMap<int>()[0]; |
| 204 | } |
| 205 | for (int i=0; i<N; ++i) { |
| 206 | outputs[i] = bodyOutputs[i+1]; |
| 207 | } |
| 208 | } |
| 209 | for (int i=0; i<K; ++i) { |
| 210 | outputs[i+N] = _Stack(spans[i]); |
| 211 | } |
| 212 | return outputs; |
| 213 | } |
| 214 | std::vector<Express::VARP> condInputs(mInfo->mCondInputNumber); |
| 215 | for (auto& p : mInfo->mInputForCond) { |
| 216 | condInputs[p.first] = inputs[p.second]; |
| 217 | } |
| 218 | for (auto& p : mInfo->mInputForBody) { |
| 219 | bodyInputs[p.first] = inputs[p.second]; |
| 220 | } |
| 221 | |
| 222 | for (int i = 0; i < mInfo->mOutputFromInput.size(); ++i) { |
| 223 | outputs[i] = inputs[mInfo->mOutputFromInput[i]]; |
| 224 | } |
| 225 | while (true) { |
| 226 | VARP res; |
| 227 | { |