| 168 | |
| 169 | template<typename T, af_op_t op> |
| 170 | void sparseArithOp(Param<T> oVals, Param<int> oColIdx, CParam<int> oRowIdx, |
| 171 | const uint Rows, CParam<T> lvals, CParam<int> lRowIdx, |
| 172 | CParam<int> lColIdx, CParam<T> rvals, CParam<int> rRowIdx, |
| 173 | CParam<int> rColIdx) { |
| 174 | const int *orPtr = oRowIdx.get(); |
| 175 | const T *lvPtr = lvals.get(); |
| 176 | const int *lrPtr = lRowIdx.get(); |
| 177 | const int *lcPtr = lColIdx.get(); |
| 178 | const T *rvPtr = rvals.get(); |
| 179 | const int *rrPtr = rRowIdx.get(); |
| 180 | const int *rcPtr = rColIdx.get(); |
| 181 | |
| 182 | arith_op<T, op> binOp; |
| 183 | |
| 184 | auto ZERO = scalar<T>(0); |
| 185 | |
| 186 | for (uint row = 0; row < Rows; ++row) { |
| 187 | const int lEnd = lrPtr[row + 1]; |
| 188 | const int rEnd = rrPtr[row + 1]; |
| 189 | const int offs = orPtr[row]; |
| 190 | |
| 191 | T *ovPtr = oVals.get() + offs; |
| 192 | int *ocPtr = oColIdx.get() + offs; |
| 193 | |
| 194 | uint rowNNZ = 0; |
| 195 | int l = lrPtr[row]; |
| 196 | int r = rrPtr[row]; |
| 197 | while (l < lEnd && r < rEnd) { |
| 198 | int lci = lcPtr[l]; |
| 199 | int rci = rcPtr[r]; |
| 200 | |
| 201 | T lhs = (lci <= rci ? lvPtr[l] : ZERO); |
| 202 | T rhs = (lci >= rci ? rvPtr[r] : ZERO); |
| 203 | |
| 204 | ovPtr[rowNNZ] = binOp(lhs, rhs); |
| 205 | ocPtr[rowNNZ] = (lci <= rci) ? lci : rci; |
| 206 | |
| 207 | l += (lci <= rci); |
| 208 | r += (lci >= rci); |
| 209 | rowNNZ++; |
| 210 | } |
| 211 | while (l < lEnd) { |
| 212 | ovPtr[rowNNZ] = binOp(lvPtr[l], ZERO); |
| 213 | ocPtr[rowNNZ] = lcPtr[l]; |
| 214 | l++; |
| 215 | rowNNZ++; |
| 216 | } |
| 217 | while (r < rEnd) { |
| 218 | ovPtr[rowNNZ] = binOp(ZERO, rvPtr[r]); |
| 219 | ocPtr[rowNNZ] = rcPtr[r]; |
| 220 | r++; |
| 221 | rowNNZ++; |
| 222 | } |
| 223 | } |
| 224 | } |
| 225 | } // namespace kernel |
| 226 | } // namespace cpu |
| 227 | } // namespace arrayfire |