| 136 | } |
| 137 | |
| 138 | __MATX_INLINE__ auto get_jit_op_str() const { |
| 139 | std::string func_name = get_jit_class_name(); |
| 140 | cuda::std::array<index_t, RANK + 1> out_dims_; |
| 141 | for (int i = 0; i < RANK + 1; i++) { |
| 142 | out_dims_[i] = Size(i); |
| 143 | } |
| 144 | |
| 145 | return cuda::std::make_tuple( |
| 146 | func_name, |
| 147 | std::format("template <{}> struct {} {{\n" |
| 148 | " using value_type = typename T0::value_type;\n" |
| 149 | " using matxop = bool;\n" |
| 150 | " constexpr static int RANK_ = {};\n" |
| 151 | " constexpr static cuda::std::array<index_t, RANK_+1> sizes_ = {{ {} }};\n" |
| 152 | " constexpr static int axis_ = {};\n" |
| 153 | " {}" |
| 154 | " // Const GetVal\n" |
| 155 | " template <typename CapType, int I, int N>\n" |
| 156 | " __MATX_INLINE__ __MATX_DEVICE__ auto GetVal(index_t oidx, cuda::std::array<index_t, RANK_>& indices) const {{\n" |
| 157 | " if constexpr ( I == N ) {{\n" |
| 158 | " const auto &op = cuda::std::get<0>(ops_);\n" |
| 159 | " return get_value<CapType>(op, indices);\n" |
| 160 | " }} else {{\n" |
| 161 | " if ( I < oidx ) {{\n" |
| 162 | " return GetVal<CapType, I+1, N>(oidx, indices);\n" |
| 163 | " }} else {{\n" |
| 164 | " const auto &op = cuda::std::get<I>(ops_);\n" |
| 165 | " return get_value<CapType>(op, indices);\n" |
| 166 | " }}\n" |
| 167 | " }}\n" |
| 168 | " }}\n" |
| 169 | " // Non-const GetVal for lvalue assignments\n" |
| 170 | " template <typename CapType, int I, int N>\n" |
| 171 | " __MATX_INLINE__ __MATX_DEVICE__ decltype(auto) GetVal(index_t oidx, cuda::std::array<index_t, RANK_>& indices) {{\n" |
| 172 | " if constexpr ( I == N ) {{\n" |
| 173 | " auto &op = cuda::std::get<0>(ops_);\n" |
| 174 | " return get_value<CapType>(op, indices);\n" |
| 175 | " }} else {{\n" |
| 176 | " if ( I < oidx ) {{\n" |
| 177 | " return GetVal<CapType, I+1, N>(oidx, indices);\n" |
| 178 | " }} else {{\n" |
| 179 | " auto &op = cuda::std::get<I>(ops_);\n" |
| 180 | " return get_value<CapType>(op, indices);\n" |
| 181 | " }}\n" |
| 182 | " }}\n" |
| 183 | " }}\n" |
| 184 | " // Const operator()\n" |
| 185 | " template <typename CapType, typename... Is>\n" |
| 186 | " __MATX_INLINE__ __MATX_DEVICE__ auto operator()(Is... is) const {{\n" |
| 187 | " if constexpr (CapType::ept == ElementsPerThread::ONE) {{\n" |
| 188 | " cuda::std::array<index_t, RANK_+1> indices{{is...}};\n" |
| 189 | " cuda::std::array<index_t, RANK_> indices_o;\n" |
| 190 | " index_t oidx = indices[axis_];\n" |
| 191 | " for(int i = 0; i < axis_; i++) {{\n" |
| 192 | " indices_o[i] = indices[i];\n" |
| 193 | " }}\n" |
| 194 | " for(int i = axis_; i < (int)indices_o.size(); i++) {{\n" |
| 195 | " indices_o[i] = indices[i+1];\n" |