Skip to content
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.

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

implement batched serial iamax #2399

Merged
merged 5 commits into from
Oct 29, 2024

Conversation

yasahi-hpc
Copy link
Contributor

This PR implements iamax function, which is needed for getrf PR.

Following files are added:

  1. KokkosBatched_Iamax_Serial_Impl.hpp: Internal interfaces with implementation details
  2. KokkosBatched_Iamax.hpp: APIs
  3. Test_Batched_SerialIamax.hpp: Unit tests for that

Detailed description

This returns the index of the first element having maximum absolute value.

  • X: (batch_count, n)
    The length N vector.

Parallelization would be made in the following manner. This is efficient only when
A is given in LayoutLeft for GPUs and LayoutRight for CPUs (parallelized over batch direction).

Kokkos::parallel_for('iamax', 
    Kokkos::RangePolicy<execution_space> policy(0, n),
    [=](const int k) {
        auto sub_x = Kokkos::subview(m_x, k, Kokkos::ALL());
        auto iamax = KokkosBatched::SerialIamax::invoke(sub_x);
        m_r(k)     = iamax;
    });

Tests

  1. Analytical tests for length 3 vectors
A0: [1, 2, 0] -> 1
A1: [-5, 4, 3] -> 0
A2: [0, 0, 0] -> 0
A3: [0, -1, -1] -> 1
  1. Make a random vector from random X. Compare the return value with the index of the first element having maximum absolute value.

Copy link
Contributor

@lucbv lucbv left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks fine, the only small question would be: "Are we worried about views containing more than 2B elements with an int return type?"

@yasahi-hpc
Copy link
Contributor Author

yasahi-hpc commented Oct 28, 2024

Looks fine, the only small question would be: "Are we worried about views containing more than 2B elements with an int return type?"

Assuming that this is a KokkosBatched functionality, return value should be smaller than 2B elements for a practical use case. However, if you are thinking of security issue like #2397, it makes sense to use int64_t. If you prefer, I can fix it accordingly.

@lucbv
Copy link
Contributor

lucbv commented Oct 28, 2024

I am not really worried was just thinking about it. One thing we could do is check the type that the view is using to store indices and use that. Then the problem becomes a Kokkos Core problem :p

@yasahi-hpc
Copy link
Contributor Author

Sure.
I will rely on the index_type of a View.

Yuuichi Asahi added 5 commits October 29, 2024 04:42
@yasahi-hpc yasahi-hpc force-pushed the implement-batched-serial-iamax branch from c0ccdb9 to b0ab297 Compare October 28, 2024 19:43
Copy link
Contributor

@lucbv lucbv left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks pretty good to me!

@lucbv lucbv merged commit e1265ec into kokkos:develop Oct 29, 2024
19 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants