New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve batched serial trsm implementation and testing #2432

Open

yasahi-hpc wants to merge 8 commits into kokkos:develop from yasahi-hpc:improve-batched-serial-trsm

+1,494 −300

Contributor

yasahi-hpc commented Nov 18, 2024 •

edited

Loading

This PR aims at improving the implementation and testing of serial Trsm.

Moving implementation details into Impl namespace.
Add all the specializations of the combinations of Left/Right, Upper/Lower, Non-Trans/Trans/ConjTrans, and Unit/Non-Unit
Disallow to use this function if X is a rank 1 View. Use Trsv for this case
Covering all the unit-tests for all the combinations of Left/Right, Upper/Lower, Non-Trans/Trans/ConjTrans, and Unit/Non-Unit

As a TO DO task, we need to add a ConjTrans implementation of blocked version which requires a little more investigation

Edited 25/Nov

Allow again to use this function if X is a rank 1 View

cwpearson added the AT2-CI-APPROVAL label

Contributor

cwpearson commented Nov 25, 2024 •

edited

Loading

@yasahi-hpc I'm using your PR to test a change to our CI infrastructure (related to this AT2-CI-APPROVAL label I added). Sorry if there is some noise.

Contributor Author

yasahi-hpc commented Nov 25, 2024

@yasahi-hpc I'm using your PR to test a change to our CI infrastructure (related to this AT2-CI-APPROVAL label I added). Sorry if there is some noise.

Hi @cwpearson
Please feel free to try
It would be rather appreciated if I could get some GPU runs

Contributor

cwpearson commented Nov 26, 2024

@yasahi-hpc looks like a legit failure in github-AT2 / spr / PR_SPR_ONEAPI202310_OPENMP_LEFT_MKLBLAS_MKLLAPACK_REL.

yasahi-hpc force-pushed the improve-batched-serial-trsm branch from d7da4c6 to 5950448 Compare

November 27, 2024 12:59

Contributor Author

yasahi-hpc commented Nov 27, 2024

@yasahi-hpc looks like a legit failure in github-AT2 / spr / PR_SPR_ONEAPI202310_OPENMP_LEFT_MKLBLAS_MKLLAPACK_REL.

Sorry for the mistake. It should be fine now. For some reason, I need to relax the tolerance for Intel CPU build.

lucbv requested changes

View reviewed changes

Contributor

lucbv left a comment

Some clean-ups required to simplify the code and future maintenance.

batched/dense/impl/KokkosBatched_Trsm_Serial_Internal.hpp Outdated

+              #pragma unroll
+              #endif
+                        for (int j = 0; j < jend; ++j)
+                          B2[i * bs0 + j * bs1] -= Kokkos::ArithTraits<ValueType>::conj(a21[i * as0]) * b1t[j * bs1];

Contributor

lucbv Dec 3, 2024

Suggested change

      
                        B2[i * bs0 + j * bs1] -= Kokkos::ArithTraits<ValueType>::conj(a21[i * as0]) * b1t[j * bs1];
          
                        B2[i * bs0 + j * bs1] -= (do_conj ? Kokkos::ArithTraits<ValueType>::conj(a21[i * as0]) * b1t[j * bs1] : a21[i * as0] * b1t[j * bs1]);

With this you can remove a lot of the code duplication introduce for the conjugate type.

Contributor

lucbv Dec 3, 2024

After checking Kokkos::ArithTraits, since we define it to be a no-op for non-complex floating points, you can just call on both complex and non-complex numbers, it will do the right thing.

batched/dense/impl/KokkosBatched_Trsm_Serial_Internal.hpp Outdated

               template <>
               template <typename ScalarType, typename ValueType>
               KOKKOS_INLINE_FUNCTION int SerialTrsmInternalLeftLower<Algo::Trsm::Blocked>::invoke(
-                  const bool use_unit_diag, const int m, const int n, const ScalarType alpha, const ValueType *KOKKOS_RESTRICT A,
-                  const int as0, const int as1,
+                  const bool use_unit_diag, [[maybe_unused]] const bool do_conj, const int m, const int n, const ScalarType alpha,

Contributor

lucbv Dec 3, 2024

Suggested change

      
                const bool use_unit_diag, [[maybe_unused]] const bool do_conj, const int m, const int n, const ScalarType alpha,
          
                const bool use_unit_diag, const bool /* do_conj */, const int m, const int n, const ScalarType alpha,

It looks more like it is never used... in that case just omit the variable name in the signature.

batched/dense/impl/KokkosBatched_Trsm_Serial_Internal.hpp Outdated

    
                    if (!use_unit_diag) {

                      const ValueType alpha11 = A[p * as0 + p * as1];

                      if (!use_unit_diag) {

                        const ValueType alpha11 = Kokkos::ArithTraits<ValueType>::conj(A[p * as0 + p * as1]);

Contributor

lucbv Dec 3, 2024

Suggested change

      
                      const ValueType alpha11 = Kokkos::ArithTraits<ValueType>::conj(A[p * as0 + p * as1]);
          
                      const ValueType alpha11 = (do_conj ? Kokkos::ArithTraits<ValueType>::conj(A[p * as0 + p * as1]) : A[p * as0 + p * as1]);

Simplify the code.

batched/dense/impl/KokkosBatched_Trsm_Serial_Internal.hpp Outdated

-                      const ValueType alpha11 = A[p * as0 + p * as1];
+                    if (do_conj) {
+                      if (!use_unit_diag) {
+                        const ValueType alpha11 = Kokkos::ArithTraits<ValueType>::conj(A[p * as0 + p * as1]);

Contributor

lucbv Dec 3, 2024

Same as above

batched/dense/impl/KokkosBatched_Trsm_Serial_Internal.hpp Outdated

               #if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
               #pragma unroll
               #endif
-                        for (int j = 0; j < jend; ++j) B0[i * bs0 + j * bs1] -= a01[i * as0] * b1t[j * bs1];
+                          for (int j = 0; j < jend; ++j)
+                            B0[i * bs0 + j * bs1] -= Kokkos::ArithTraits<ValueType>::conj(a01[i * as0]) * b1t[j * bs1];

Contributor

lucbv Dec 3, 2024

Same as above

batched/dense/impl/KokkosBatched_Trsm_Serial_Internal.hpp Outdated

               template <>
               template <typename ScalarType, typename ValueType>
               KOKKOS_INLINE_FUNCTION int SerialTrsmInternalLeftUpper<Algo::Trsm::Blocked>::invoke(
-                  const bool use_unit_diag, const int m, const int n, const ScalarType alpha, const ValueType *KOKKOS_RESTRICT A,
-                  const int as0, const int as1,
+                  const bool use_unit_diag, [[maybe_unused]] const bool do_conj, const int m, const int n, const ScalarType alpha,

Contributor

lucbv Dec 3, 2024

Same as above

cwpearson added AT2-CI-APPROVAL and removed AT2-CI-APPROVAL labels

Yuuichi Asahi added 8 commits

December 4, 2024 03:11


          Use Trsv instead of Trsm if X is a rank 1 matrix

9ad7ddd

Signed-off-by: Yuuichi Asahi <[email protected]>


          Add missing specialization of Trsm serial implementation

eec6cbf

Signed-off-by: Yuuichi Asahi <[email protected]>


          Add missing tests for serial Trsm

2181aa3

Signed-off-by: Yuuichi Asahi <[email protected]>


          fix: format

dd4de1c

Signed-off-by: Yuuichi Asahi <[email protected]>


          fix: initialization order based on codeQL

a2e34b2

Signed-off-by: Yuuichi Asahi <[email protected]>


          Allow trsm serial to work on rank 1 matrix

c472297

Signed-off-by: Yuuichi Asahi <[email protected]>


          fix: MKL interface of Serial Trsm

dbd8e60

Signed-off-by: Yuuichi Asahi <[email protected]>


          simplify do_conj logic in trsm serial internal

af95556

Signed-off-by: Yuuichi Asahi <[email protected]>

yasahi-hpc force-pushed the improve-batched-serial-trsm branch from 5950448 to af95556 Compare

December 3, 2024 18:12

yasahi-hpc requested a review from lucbv

December 3, 2024 19:09

cwpearson added AT2-CI-APPROVAL and removed AT2-CI-APPROVAL labels

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

AT2-CI-APPROVAL