KokkosBlas: limited vector size in norm calculation?

[kddevin]

Bug Report

@trilinos/tpetra

Description

@jennloe reports

I have this Trilinos (Belos + Tpetra) code which is crashing, and I don't understand why

76   typedef double                            ScalarType;
77   typedef int                               OT;
78   typedef Tpetra::MultiVector<ScalarType>   MV;
79   typedef Tpetra::MultiVector<ScalarType>::mag_type   MT;
...
132   RCP<Tpetra::MultiVector<ScalarType>> X1 = rcp( new Tpetra::MultiVector<ScalarType>(map, 1008) );
133   X1->randomize();
134   RCP<Tpetra::MultiVector<ScalarType>> XCopy = rcp(new Tpetra::MultiVector<ScalarType>(*X1, Teuchos::Copy)); //Deep copy of X1.
....
154   // DEBUG: //Verify coefficients:
155   MV coeffs_mv = makeStaticLocalMultiVector (*X1, 1008, 1008);
156   std::cout << "Starting deep copy" << std::endl;
157   Tpetra::deep_copy(coeffs_mv, *coeffMat);
158   std::cout << "Finished deep copy" << std::endl;
159   XCopy->multiply(Teuchos::NO_TRANS, Teuchos::NO_TRANS, 1.0, *X1, coeffs_mv, -1.0);
160   std::cout << "Passed multiply." << std::endl;
161   std::vector<MT> norms(1008);
162   Teuchos::ArrayView<MT> normView(norms);
163   std::cout << "created norms array" << std::endl;
164   XCopy->norm2(normView);
165   std::cout << "Here are the QR-Orig norms.  Should be zero." << std::endl;
166   std::cout << normView << std::endl << std::endl;

And the output is:

$ ./Belos_TpetraOrthogCoeffs.exe
Starting deep copy
Finished deep copy
Passed multiply.
created norms array
terminate called after throwing an instance of 'std::runtime_error'
  what():  cudaStreamSynchronize(m_stream) error( cudaErrorIllegalAddress): an illegal memory access was encountered /ascldap/sers/jloe/Trilinos/packages/kokkos/core/src/Cuda/Kokkos_Cuda_Instance.cpp:314
Traceback functionality not available
Aborted (core dumped)

At first glance, the problem appears to be a mismatch between the norms argument (on host) and the vector argument (on device)

Steps to Reproduce

Reproduced with simple test program on ascicgpu with no UVM.

[kddevin]

I was too quick to say it was a Tpetra issue. The problem seems to be in KokkosBlas.

I wrote the following reproducer:

  const size_t nGlobalEntries = 100;
  const size_t nVecs = 181;

  Kokkos::View<Scalar **, Kokkos::LayoutLeft, typename Node::device_type> mv("mv", nGlobalEntries, nVecs);
  Kokkos::deep_copy(mv, 3.14);

  std::vector<Scalar> norm(nVecs);
  Kokkos::View<Scalar*, Kokkos::HostSpace> normView(&norm[0], nVecs);
  try {
    KokkosBlas::nrm2_squared(normView, mv);
  }
  catch (std::exception &e) {
    std::cout << "Test failed for " << nVecs << " vectors:  " << e.what() << std::endl;
    TEST_ASSERT(false);
  }

With Node::device_type == Kokkos::Cuda, this code fails with the error that @jennloe saw.
If I change only nVecs from 181 to 180, this code passes.
@trilinos/kokkos-kernels

[kddevin]

Brian says this is a bug with Kokkos' reduce; see kokkos/kokkos#4461

[brian-kelley]

@kddevin @jennloe FYI kokkos/kokkos-kernels#1204 will take care of this. You can have any number of vectors on any backend, and for most of the BLAS-1 functions (dot, norms other than nrminf, sum) you'll get a ~3x speedup on Cuda compared to before.

[jhux2]

@iyamazaki