Method MultTransAdd

ngscuda/cuda_linalg.cpp:538–600

  

  void DevConstantElementByElementMatrix ::
  MultTransAdd (double s, const BaseVector & x, BaseVector & y) const
  {
    static Timer t("DevConstantEBEMatrix::MultTransAdd"); RegionTimer reg(t);
    
    UnifiedVectorWrapper ux(x);
    UnifiedVectorWrapper uy(y);
    
    ux.UpdateDevice();
    uy.UpdateDevice();
    
    auto hm = devmat.Height();
    auto wm = devmat.Width();
    // if (disjoint_rows)
    if (true)      
      {
        DevStackArray<double> dev_hx(numblocks*hm);
        DevStackArray<double> dev_hy(numblocks*wm);

        // ConstEBEKernelCopyIn (numblocks, hm, coldnums.DevData(), (double*)ux.DevData(), dev_hx.DevData());
        DeviceParallelFor
          (numblocks*hm,
           [locx=dev_hx.DevData(), globx=ux.DevData(), idx=coldnums.DevData()] DEVICE_LAMBDA (auto tid)
           {
             locx[tid] = globx[idx[tid]];
           });
        
        // dev_hy = dev_hx * mat
        
        FlatMatrix<Dev<double>> matx(numblocks, hm, dev_hx.Data());
        FlatMatrix<Dev<double>> maty(numblocks, wm, dev_hy.Data());
        MultMatMat (matx, devmat, maty, s, 0);
        
        // ConstEBEKernelCopyOut (numblocks, wm, rowdnums.DevData(), dev_hy.DevData(), (double*)uy.DevData());
        DeviceParallelFor
          (numblocks*wm,
           [globy=uy.DevData(), locy=dev_hy.DevData(), idx=rowdnums.DevData()] DEVICE_LAMBDA (auto tid)
           {
             atomicAdd((double*)globy+idx[tid], locy[tid]);
           });
        
      }
    else
      {
        for (auto c : row_coloring)
          {
            DevStackArray<double> dev_hx(c.Size()*hm);
            DevStackArray<double> dev_hy(c.Size()*wm);

            ConstEBEKernelCopyInIdx (c.Size(), (int*)c.Data(), hm, coldnums.DevData(), (double*)ux.DevData(), dev_hx.DevData());
            // dev_hy = dev_hx * mat

            FlatMatrix<Dev<double>> matx(c.Size(), hm, dev_hx.Data());
            FlatMatrix<Dev<double>> maty(c.Size(), wm, dev_hy.Data());
            MultMatMat (matx, devmat, maty, s, 0);
           
            ConstEBEKernelCopyOutIdx (c.Size(), (int*)c.Data(), wm, rowdnums.DevData(), dev_hy.DevData(), (double*)uy.DevData());
          }