Skip to content

Conversation

@horizon-blue
Copy link
Contributor

@horizon-blue horizon-blue commented Nov 20, 2025

(Closes MET-47)

Summary of Changes

This PR addresses some of the suggestions that @mugamma brought up in #14. In particular, it defines the operator[] on Array2D to return a 1D view of a row, so we can use patterns like array2d[i][j] instead of array2d(i, j) to access the element.

Another nice thing about returning the 1D span is that we can use range-based for loop to go over the elements in a row as well, e.g.

for (auto& val : array2d[rows]) { /* do something with val*/ }

You can find some example usages in the included test file.

Another minor change in this PR is the refactoring ofArray2D methods to use the new CUDA_CALLABLE macro that Arijit introduced recently.

Test Plan

To run the included unit tests:

pixi run test

@linear
Copy link

linear bot commented Nov 20, 2025

MET-47 Improve `Array2D` semantics

To address some of the comments in #14

@horizon-blue horizon-blue marked this pull request as ready for review November 20, 2025 23:53
@arijit-dasgupta
Copy link
Contributor

arijit-dasgupta commented Nov 22, 2025

@horizon-blue I leave this to you, but I did prompt cursor for what kind of interesting tests we could do with Array2D, and it suggested a barrage of things I did not like ("boo LLMs"), but it did bring up a couple I like. Thought of sharing them with you here:

// Test that modifications through view affect underlying data
TYPED_TEST(Array2DTestFixture, ViewModifiesUnderlyingData) {
    if constexpr (std::is_same_v<TypeParam, std::vector<float>>) {
        uint32_t rows = 3;
        uint32_t cols = 4;
        auto data = TypeParam(rows * cols, 0.0f);
        auto array2d = Array2D(thrust::raw_pointer_cast(data.data()), rows, cols);

        // Modify through view
        array2d[1][2] = 42.5f;
        // Verify underlying data changed
        EXPECT_FLOAT_EQ(data[1 * cols + 2], 42.5f);

        // Modify underlying data directly
        data[0 * cols + 1] = 99.9f;
        // Verify view reflects change
        EXPECT_FLOAT_EQ(array2d[0][1], 99.9f);
    }
}

// Test multiple views of the same data
TYPED_TEST(Array2DTestFixture, MultipleViewsOfSameData) {
    if constexpr (std::is_same_v<TypeParam, std::vector<float>>) {
        uint32_t rows = 2;
        uint32_t cols = 3;
        auto data = TypeParam(rows * cols, 0.0f);
        auto view1 = Array2D(thrust::raw_pointer_cast(data.data()), rows, cols);
        auto view2 = Array2D(thrust::raw_pointer_cast(data.data()), rows, cols);

        // Modify through view1
        view1[0][0] = 100.0f;
        // Verify view2 sees the change
        EXPECT_FLOAT_EQ(view2[0][0], 100.0f);

        // Modify through view2
        view2[1][2] = 200.0f;
        // Verify view1 sees the change
        EXPECT_FLOAT_EQ(view1[1][2], 200.0f);
    }
}

Considering that this is a view on an underlying data, these tests made sense to me. For ViewModifiesUnderlyingData, it's good to verify that any underlying data modification updates the view, and that editing via the view also changes the underlying data. For MultipleViewsOfSameData, it is helpful to verify that one view of the data ultimately looks at the same underlying data even if another view modified it before. This reminded me of how shared memory in NVIDIA GPUs work.

I leave it to you to decide if you want to add them in, they pass on my side locally so either ways this works for sure!

Otherwise nice change with the operator! Will be super convenient. I'll approve this, and you can merge it when you decide what to do!

}

__host__ __device__ constexpr auto rank() const noexcept {
CUDA_CALLABLE constexpr auto rank() const noexcept {
Copy link
Contributor

Choose a reason for hiding this comment

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

Suggested change
CUDA_CALLABLE constexpr auto rank() const noexcept {
CUDA_CALLABLE constexpr auto ndim() const noexcept {

Copy link
Contributor

Choose a reason for hiding this comment

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

Extremely minor thing, I know this rank here refers to the dynamic rank of a matrix, which is 2, but I tend to think of the linear algebra rank instead. suggestion is to rename it to ndim (copying pythonic ndarrays), but feel free to ignore haha.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Sounds good to me haha. I named the method rank() only because the mdspan is calling it so, but I agree that this can be confusing for folks with a math brain 😆. I will update this shortly

Copy link
Contributor

Choose a reason for hiding this comment

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

If you accept my ndim suggestion, then you would need to edit line 147 from rank to ndim for test to pass. It works on my end locally.

@horizon-blue horizon-blue force-pushed the xiaoyan/array2d-improvements branch from c611bbb to 96ce9ba Compare November 23, 2025 05:28
@horizon-blue
Copy link
Contributor Author

Ok, I renamed rank() to ndim(), added the additional tests that @arijit-dasgupta suggested, and modified constructor slightly so we don't have to explicitly invoke raw_pointer_cast when instantiating Array2D. I'm going to merge this PR now to reduce the amount of pending changes, but feel free to leave comment if you're against any of the changes and I will address them later :)

@horizon-blue horizon-blue merged commit 4bee915 into master Nov 23, 2025
1 check passed
@horizon-blue horizon-blue deleted the xiaoyan/array2d-improvements branch November 23, 2025 05:36
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.

3 participants