| 1024 | } |
| 1025 | |
| 1026 | std::vector<array> compile_replace( |
| 1027 | const std::vector<array>& tape, |
| 1028 | const std::vector<array>& trace_inputs, |
| 1029 | const std::vector<array>& trace_outputs, |
| 1030 | const std::vector<array>& inputs, |
| 1031 | bool shapeless) { |
| 1032 | std::unordered_map<uintptr_t, array> trace_to_real; |
| 1033 | for (int i = 0; i < inputs.size(); ++i) { |
| 1034 | trace_to_real.insert({trace_inputs[i].id(), inputs[i]}); |
| 1035 | } |
| 1036 | |
| 1037 | auto is_load = [](const Primitive& p) { return typeid(p) == typeid(Load); }; |
| 1038 | |
| 1039 | for (auto& a : tape) { |
| 1040 | // Arrays in the tape without primitives are either: |
| 1041 | // - inputs, which are already in the map |
| 1042 | // - constants, which can be used directly |
| 1043 | // - a load primitive which has no inputs and will become a constant |
| 1044 | // after the first eval |
| 1045 | if (!a.has_primitive() || is_load(a.primitive())) { |
| 1046 | trace_to_real.insert({a.id(), a}); |
| 1047 | } else { |
| 1048 | // Find real inputs |
| 1049 | std::vector<array> real_inputs; |
| 1050 | for (auto& in : a.inputs()) { |
| 1051 | real_inputs.push_back(trace_to_real.at(in.id())); |
| 1052 | } |
| 1053 | if (a.siblings().empty()) { |
| 1054 | auto shape = |
| 1055 | shapeless ? a.primitive().output_shapes(real_inputs)[0] : a.shape(); |
| 1056 | auto real_a = array( |
| 1057 | std::move(shape), |
| 1058 | a.dtype(), |
| 1059 | a.primitive_ptr(), |
| 1060 | std::move(real_inputs)); |
| 1061 | trace_to_real.insert({a.id(), std::move(real_a)}); |
| 1062 | } else { |
| 1063 | // Ensure the order is correct for multi-output primitives |
| 1064 | std::vector<Dtype> types; |
| 1065 | auto trace_out = a.outputs(); |
| 1066 | for (auto& o : trace_out) { |
| 1067 | types.push_back(o.dtype()); |
| 1068 | } |
| 1069 | std::vector<Shape> shapes; |
| 1070 | if (shapeless) { |
| 1071 | shapes = a.primitive().output_shapes(real_inputs); |
| 1072 | } else { |
| 1073 | for (auto& o : trace_out) { |
| 1074 | shapes.push_back(o.shape()); |
| 1075 | } |
| 1076 | } |
| 1077 | auto real_out = array::make_arrays( |
| 1078 | std::move(shapes), types, a.primitive_ptr(), real_inputs); |
| 1079 | for (int i = 0; i < trace_out.size(); ++i) { |
| 1080 | trace_to_real.insert({trace_out[i].id(), std::move(real_out[i])}); |
| 1081 | } |
| 1082 | } |
| 1083 | } |