Add flat sparse packs #877

Yurlungur · 2023-05-22T20:39:58Z

PR Summary

In riot I recently encountered an interesting use-case, which I needed to support. It's actually familiar. In the boundary communication kernels, we fuse the loop over meshblocks, over faces, and over variables. So the outermost index has a range that spans ALL variables over ALL boundary buffers and ALL meshblocks in a given MeshData object. We do this so that when we run with hierarchical parallelism, there's enough "chunks" of work in the outermost index space to saturate the device... i.e., we need at least as many elements in the outer range as we have warps on the GPU.

I found myself needing the same capability for loops over variables in the interior of a block. I.e., I needed a hierarchical structure where the outer loop is over meshblocks and variables TOGETHER and the inner loop is over the 3D index space of the meshblock. Because of sparse, this outer loop needs to be a SINGLE loop. The number of variables on each meshblock might vary.

Here I implement this capability in @lroberts36 's SparsePacks, which was the easiest place. When built with the sparse flag in the generic factory function, or with the GetFlat factory function, the sparse pack builds with this shape.

Note that I added a blurb describing this functionality, but the documentation for the SparsePack is very sparse, and I didn't fill it in.

PR Checklist

Code passes cpplint
New features are documented.
Adds a test for any bugs fixed. Adds tests for new features.
Code is formatted
Changes are summarized in CHANGELOG.md
CI has been triggered on Darwin for performance regression tests.
Docs build
(@lanl.gov employees) Update copyright on changed files

Yurlungur · 2023-05-22T20:41:09Z

P.S., thanks @lroberts36 for your help with this

src/interface/sparse_pack_base.cpp

tst/unit/test_sparse_pack.cpp

brryan

LGTM, for my own understanding, how is this intended to be used downstream with sparse variables? Is it just an if check for each variable on each meshblock and the work is skipped if that variable isnt allocated on that meshblock?

src/interface/sparse_pack.hpp

src/interface/sparse_pack_base.hpp

Yurlungur · 2023-05-22T21:00:19Z

how is this intended to be used downstream with sparse variables? Is it just an if check for each variable on each meshblock and the work is skipped if that variable isnt allocated on that meshblock?

It only packs allocated variables. So if for example, I have a total of 20 variables, but 3 variables are allocated on each of 5 blocks (where the set of variables allocated are different on each block), you get an outer index of length 15.

lroberts36 · 2023-05-23T17:23:21Z

the documentation for the SparsePack is very sparse

I see what you did there ;)

Yurlungur · 2023-05-23T17:50:01Z

@lroberts36 you should probably be the other person to review this---seeing as you wrote sparse packs.

lroberts36

This LGTM except that I think we should make sure the pack bounds arrays are meaningful even for flat packs.

lroberts36 · 2023-05-23T17:30:39Z

src/interface/sparse_pack.hpp

@@ -200,27 +252,32 @@ class SparsePack : public SparsePackBase {
  KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b) const { return 0; }

  KOKKOS_INLINE_FUNCTION int GetUpperBound(const int b) const {
+    PARTHENON_DEBUG_REQUIRE(!flat_, "Bounds not valid for flat packs");


Are they not valid if b==0? I thought they were.

Ok, after reading below I see they are not valid. I think we should make it so that they are valid, and maybe add a default const int b=0 here so that GetLowerBound() and GetUpperBound() work for sparse packs (or maybe make them separate functions that have a check if flat_==true). I think this is clearer than exposing the pack size through GetMaxNumberOfVars().

Maybe better to just call them GetFlatLowerBound() and GetFlatUpperBound().

I'm going to make them valid and just re-use the API. I'm not sure if we should come up with new or more generic names, I just chose not to for now. We can revisit that later maybe.

src/interface/sparse_pack_base.cpp

lroberts36 · 2023-05-23T17:41:39Z

src/interface/sparse_pack_base.cpp

+      if (!desc.flat) {
+        pack.bounds_h_(0, b, i) = idx;
+      }


Can't we still set this with the actual block index, i.e. pack.bounds_h_(0, block, i) = idx? I am not sure what the use case would be, but it would allow all of the Get*Bound stuff to work returning the flat index.

Is this actually desired behaviour though? I wonder if GetBounds should return 0, and nvars always in the flat case? I could add an if statement to that effect maybe?

Now updated GetLowerBound returns 0 and GetUpperBound returns total numbers of vars and blocks (i.e., pack size) for flat packs.

src/interface/sparse_pack_base.cpp

pgrete

Looks good to me.
I don't have a strong opinion with respect to the index handling.
In the spirit of our ParArrayNDs that transparently handle 0 I could see the appeal to also allow using the standard () operator and bounds being used with a PARTHENON_DEBUG_REQUIRE(!flat || b == 0,...).

doc/sphinx/src/interface/sparse.rst

src/interface/sparse_pack.hpp

Co-authored-by: Philipp Grete <pgrete@hs.uni-hamburg.de>

Yurlungur · 2023-05-26T00:13:30Z

After incorporating @lroberts36 's suggestions, I broke a test, so I'm going to need a little time to debug. We should not wait for this PR for the next release. @pgrete .

…flat-sparse

Yurlungur · 2023-05-26T18:41:18Z

@lroberts36 this is ready for re-review. I believe I've incorporated all your comments.

src/interface/sparse_pack.hpp

lroberts36

I think we were on a different page about how to keep the (block, variable) indexing working for flat packs. I included some possible changes that are along the lines of what I was thinking. I think we probably don't want to do it because of the extra H2D copy, but another array could be added that maps from flat idx to block idx (along with functions for querying the pack for which block an flat idx corresponds to) so that e.g. the correct coordinates could be accessed from a flat pack.

lroberts36 · 2023-05-30T17:30:20Z

src/interface/sparse_pack.hpp

-  KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b) const { return 0; }
+  KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b = 0) const { return 0; }


I think this should change for flat packs. GetLowerBound(b) with b > 0 should return bounds_(0, b, 0) and GetUpperBound(b) should return bounds_(1, b, nvar_).

lroberts36 · 2023-05-30T17:35:08Z

src/interface/sparse_pack.hpp

+  KOKKOS_INLINE_FUNCTION int GetUpperBound(const int b = 0) const {
+    return bounds_(1, !flat_ * b, !flat_ * nvar_);


I don't think there should be any check for if the pack is flat here. This is why I think that there actually needs to be a new function that returns the total upper bound

KOKKOS_INLINE_FUNCTION int GetUpperBound() const { assert(flat_); return bounds_(1, GetNBlocks() - 1, nvar_); }

lroberts36 · 2023-05-30T17:35:31Z

src/interface/sparse_pack.hpp

-  KOKKOS_INLINE_FUNCTION int GetUpperBoundHost(const int b) const {
-    return bounds_h_(1, b, nvar_);
+  KOKKOS_INLINE_FUNCTION int GetUpperBoundHost(const int b = 0) const {
+    return bounds_h_(1, !flat_ * b, !flat_ * nvar_);


Also no flat check here.

lroberts36 · 2023-05-30T17:38:06Z

src/interface/sparse_pack_base.cpp

@@ -179,13 +187,12 @@ SparsePackBase SparsePackBase::Build(T *pmd, const PackDescriptor &desc) {
      }
    }
    // Record the maximum for easy access
-    pack.bounds_h_(1, b, nvar) = idx - 1;
+    pack.bounds_h_(1, !desc.flat * b, !desc.flat * nvar) = idx - 1;


Suggested change

pack.bounds_h_(1, !desc.flat * b, !desc.flat * nvar) = idx - 1;

pack.bounds_h_(1, block, nvar) = idx - 1;

lroberts36 · 2023-05-30T17:38:46Z

src/interface/sparse_pack_base.cpp

+      if (!desc.flat) {
+        pack.bounds_h_(1, b, i) = idx - 1;
+      }


Suggested change

if (!desc.flat) {

pack.bounds_h_(1, b, i) = idx - 1;

}

pack.bounds_h_(1, b, i) = idx - 1;

lroberts36 · 2023-05-30T17:40:45Z

src/interface/sparse_pack_base.cpp

+    if (!desc.flat) {
+      idx = 0;
+      b = block;
+      // JMM: This line could be unified with the coords_h line below,
+      // but it would imply unnecessary copies in the case of non-flat
+      // packs.
+      coords_h(b) = pmbd->GetBlockPointer()->coords_device;
+    }


Suggested change

if (!desc.flat) {

idx = 0;

b = block;

// JMM: This line could be unified with the coords_h line below,

// but it would imply unnecessary copies in the case of non-flat

// packs.

coords_h(b) = pmbd->GetBlockPointer()->coords_device;

}

if (!desc.flat) {

idx = 0;

b = block;

}

coords_h(block) = pmbd->GetBlockPointer()->coords_device;

lroberts36 · 2023-05-30T17:41:27Z

src/interface/sparse_pack_base.cpp

+      if (!desc.flat) {
+        pack.bounds_h_(0, b, i) = idx;
+      }


Suggested change

if (!desc.flat) {

pack.bounds_h_(0, b, i) = idx;

}

pack.bounds_h_(0, block, i) = idx;

lroberts36 · 2023-05-30T17:41:51Z

src/interface/sparse_pack_base.cpp

+                  if (desc.flat) {
+                    coords_h(idx) = pmbd->GetBlockPointer()->coords_device;
                  }


Suggested change

if (desc.flat) {

coords_h(idx) = pmbd->GetBlockPointer()->coords_device;

}

I think this is potentially confusing that for regular backs GetCoords takes the block number as an argument while for flat packs it takes the variable flat index.

I want to be able to make this call, as I think there's a clear use-case: i.e., you need coords for your stencil op. However, I would be willing to give GetCoords two different names?

lroberts36 · 2023-05-30T17:43:23Z

src/interface/sparse_pack_base.cpp

  // Size is nvar + 1 to store the maximum idx for easy access
-  pack.bounds_ = bounds_t("bounds", 2, nblocks, nvar + 1);
+  pack.bounds_ = bounds_t("bounds", 2, desc.flat ? 1 : nblocks, !desc.flat * nvar + 1);


Suggested change

pack.bounds_ = bounds_t("bounds", 2, desc.flat ? 1 : nblocks, !desc.flat * nvar + 1);

pack.bounds_ = bounds_t("bounds", 2, nblocks, nvar + 1);

lroberts36 · 2023-05-30T17:43:37Z

src/interface/sparse_pack_base.cpp

  pack.bounds_h_ = Kokkos::create_mirror_view(pack.bounds_);

-  pack.coords_ = coords_t("coords", nblocks);
+  pack.coords_ = coords_t("coords", desc.flat ? max_size : nblocks);


Suggested change

pack.coords_ = coords_t("coords", desc.flat ? max_size : nblocks);

pack.coords_ = coords_t("coords", nblocks);

Yurlungur · 2023-05-30T17:57:13Z

@lroberts36 I don't think I understand what you want. Can we talk over MM or something?

lroberts36 · 2023-05-30T17:58:18Z

yep, I am on MM

Yurlungur · 2023-05-30T20:49:17Z

@lroberts36 take a look now. I've re-enabled original functionality for GetLowerBounds and GetUpperBounds as we discussed. Also added GetSize as a transparent way to get total size of the array.

lroberts36

There are a couple more comments/suggestions below, but this LGTM and I approve

lroberts36 · 2023-05-31T17:13:23Z

src/interface/sparse_pack_base.cpp

+                  if (desc.flat) {
+                    coords_h(idx) = pmbd->GetBlockPointer()->coords_device;
                  }


I think this is potentially confusing that for regular backs GetCoords takes the block number as an argument while for flat packs it takes the variable flat index.

lroberts36 · 2023-05-31T17:14:48Z

src/interface/sparse_pack_base.hpp

+  KOKKOS_INLINE_FUNCTION
+  Real &operator()(int idx, const int k, const int j, const int i) const {
+    PARTHENON_DEBUG_REQUIRE(flat_, "Accessor only valid for flat packs");
+    return pack_(0, 0, idx)(k, j, i);
+  }


Is there a reason for providing operator() in sparse_pack_base.hpp? I was trying to confine the interface to just SparsePack.

I think this is a mistake. I'll remove it.

Yurlungur · 2023-05-31T18:14:17Z

@pgrete the cuda regression test appears to be failing, without a useful error message, but it works on my machine. I'll dig in a little bit, but I wonder if something is wrong with the CI machine?

Yurlungur · 2023-05-31T19:07:59Z

Weird problem went away. Pulling the trigger.

jonahm-LANL added 3 commits May 22, 2023 14:20

add sparse flat packs

35a8b43

formatting

183e891

docs

426effb

Yurlungur added enhancement New feature or request Kokkos Kokkos infrastructure/feature related performance labels May 22, 2023

Yurlungur requested review from lroberts36 and pgrete May 22, 2023 20:39

Yurlungur self-assigned this May 22, 2023

changelog

6a1ad2e

Yurlungur commented May 22, 2023

View reviewed changes

src/interface/sparse_pack_base.cpp Show resolved Hide resolved

tst/unit/test_sparse_pack.cpp Show resolved Hide resolved

Yurlungur requested review from forrestglines, brryan and bprather May 22, 2023 20:43

brryan approved these changes May 22, 2023

View reviewed changes

src/interface/sparse_pack.hpp Show resolved Hide resolved

src/interface/sparse_pack.hpp Outdated Show resolved Hide resolved

src/interface/sparse_pack_base.hpp Outdated Show resolved Hide resolved

brryan comments

c7a5888

Yurlungur enabled auto-merge May 22, 2023 21:08

Yurlungur disabled auto-merge May 22, 2023 21:08

Merge branch 'develop' into jmm/flat-sparse

26299ba

lroberts36 requested changes May 23, 2023

View reviewed changes

pgrete approved these changes May 25, 2023

View reviewed changes

doc/sphinx/src/interface/sparse.rst Outdated Show resolved Hide resolved

src/interface/sparse_pack.hpp Show resolved Hide resolved

Yurlungur and others added 2 commits May 25, 2023 10:10

Update doc/sphinx/src/interface/sparse.rst

7b9f2e4

Co-authored-by: Philipp Grete <pgrete@hs.uni-hamburg.de>

Merge branch 'develop' into jmm/flat-sparse

ab4287d

jonahm-LANL added 2 commits May 26, 2023 12:34

incorporate lroberts suggestions

debfb7a

Merge branch 'jmm/flat-sparse' of github.com:lanl/parthenon into jmm/…

4354c8b

…flat-sparse

Yurlungur and others added 3 commits May 26, 2023 12:35

Merge branch 'develop' into jmm/flat-sparse

752b677

fix edge case for no vars found in flat case

1ea8f5d

Merge branch 'jmm/flat-sparse' of github.com:lanl/parthenon into jmm/…

13979aa

…flat-sparse

oops cellvars are vars now

e7d4526

Yurlungur commented May 26, 2023

View reviewed changes

src/interface/sparse_pack.hpp Show resolved Hide resolved

Yurlungur removed request for bprather and forrestglines May 26, 2023 18:50

Yurlungur enabled auto-merge May 26, 2023 22:47

Yurlungur disabled auto-merge May 26, 2023 22:47

testing lower/upper bound on host, not device

a9d8fdb

lroberts36 self-requested a review May 30, 2023 17:25

lroberts36 reviewed May 30, 2023

View reviewed changes

re-enable getlowerobund getupperbound

f611cc3

lroberts36 approved these changes May 31, 2023

View reviewed changes

Yurlungur enabled auto-merge May 31, 2023 19:08

fix bug related to lower bounds

8640446

Yurlungur disabled auto-merge May 31, 2023 22:31

Yurlungur enabled auto-merge May 31, 2023 22:31

lroberts36 added 3 commits May 31, 2023 17:39

Merge branch 'develop' into jmm/flat-sparse

90fc737

Move allocated check

85296f5

format

0a7c198

lroberts36 disabled auto-merge June 1, 2023 00:44

lroberts36 enabled auto-merge (squash) June 1, 2023 00:44

lroberts36 merged commit 74c1d05 into develop Jun 1, 2023

		KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b) const { return 0; }
		KOKKOS_INLINE_FUNCTION int GetLowerBound(const int b = 0) const { return 0; }

		KOKKOS_INLINE_FUNCTION int GetUpperBound(const int b = 0) const {
		return bounds_(1, !flat_ * b, !flat_ * nvar_);

	pack.bounds_h_(1, !desc.flat * b, !desc.flat * nvar) = idx - 1;
	pack.bounds_h_(1, block, nvar) = idx - 1;

	if (desc.flat) {
	coords_h(idx) = pmbd->GetBlockPointer()->coords_device;
	}

	pack.bounds_ = bounds_t("bounds", 2, desc.flat ? 1 : nblocks, !desc.flat * nvar + 1);
	pack.bounds_ = bounds_t("bounds", 2, nblocks, nvar + 1);

	pack.coords_ = coords_t("coords", desc.flat ? max_size : nblocks);
	pack.coords_ = coords_t("coords", nblocks);

Add flat sparse packs #877

Add flat sparse packs #877

Conversation

Yurlungur commented May 22, 2023 • edited Loading

PR Summary

PR Checklist

Yurlungur commented May 22, 2023

brryan left a comment

Choose a reason for hiding this comment

Yurlungur commented May 22, 2023

lroberts36 commented May 23, 2023

Yurlungur commented May 23, 2023

lroberts36 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pgrete left a comment

Choose a reason for hiding this comment

Yurlungur commented May 26, 2023

Yurlungur commented May 26, 2023

lroberts36 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Yurlungur commented May 30, 2023

lroberts36 commented May 30, 2023

Yurlungur commented May 30, 2023

lroberts36 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Yurlungur commented May 31, 2023

Yurlungur commented May 31, 2023

Yurlungur commented May 22, 2023 •

edited

Loading