[tbennun]

A memlet serves two purposes at the same time. It produces an address, and it moves bytes. Currently, the default is copy-on-access, recomputing the offset at every access. This PR adds Memlet.access_policy, which allows to modify that. At code generation, the access policy is lowered to LibraryNodes (just like the copy node today) and can now represent concepts such as pointer incrementation (cursors), buffer resource descriptors, TMA-generated reads, or general address-generation/iterators/accessors.

[FlorianDeconinck]

Ongoing testing on GEOS' FV3 Dynamics core routine of D_SW.

Testing the LoopCursor schedule for memlets (via ScheduleLoopCursors)

In the case of single K (vertical) loop and many J I (horizontal) dual-loop where code within the kernel is heavily dependent on if condition to restrict computation further

• 35% speed up while keeping validation

Overall this looks like it's clearly superior for serial CPU code

[FlorianDeconinck]

Ongoing testing on GEOS' FV3 Dynamics core routine of D_SW.

Testing the LoopCursor schedule for memlets (via ScheduleLoopCursors)

In the case of single K (vertical) loop and many J I (horizontal) dual-loop where code within the kernel is heavily dependent on if condition to restrict computation further

• 35% speed up while keeping validation

Overall this looks like it's clearly superior for serial CPU code

After more thorough look, ForLoop already to direct referencing per For axis, which seems to lead to a similar speed up. The transform remains valid (as in it validates numerics) but it adds a large amount of pointer/integer that bloats the asm

[tbennun]

cscs-ci run

[DJAntivenom]

Just as an FYI: As part of my master's thesis I'm working on TMA support and I will be adding the corresponding access policy :)

[ThrudPrimrose]

We should keep type hints always, block: CodeBlock and return type would be list[CodeBlock] ?

[ThrudPrimrose]

I think having both versions' generated code (copy-on-access and ptr-increment) would nice for documentation I think

[ThrudPrimrose]

I think this is the 3rd or 4th PR where we need to move around the horrible default GPU_SCHEDULES type which is just GPU_ThreadBlock, which is massively confusing. (Including new gpu codegen and new gpu offloading pass).

I think we should update the dtypes to have GPU_KERNEL_SCHEDULES in dtypes.py, and to replace GPU_SCHEDULES with a better name, such as GPU_LANE_SCHEDULES.

To avoid an accidental conflict I propose not using GPU_SCHEDULES but GPU_KERNEL_SCHEDULES and GPU_LANE_SCHEDULES

[ThrudPrimrose]

Would be nice addition to dtypes.py or a similar file.

[ThrudPrimrose]

Since we removed Indices, this function should only receive Subsets at this stage?

[ThrudPrimrose]

Could we achieve the same thing by querying Subset.is_contiguous_subset(array) in a cond check, and if true return subset.num_elements()?

Seem equivalent to me

[ThrudPrimrose]

Should we include maps that are sequential here?

[ThrudPrimrose]

How does this handle the nested GPU device maps?

Because from what I have seen something like this happens:

for i in [0:N:1]: at GPU Device
  for j1 in [0:M1:1]: at GPU Device
     body1
  for j2 in [0:M2:1]: at GPU Device
     body2

Then codegen implicitly generates a map that is like:

for i, j in [0:N:1, 0:max(M1,M2):1]: at GPU Device
  if j <= M1:
     body1
  if j <= M2:
     body2

So I think while this check works for cases having tblock map and not, it will have issues when we have nested gpu device maps?

To mention: #2534 (#2534 ), this pass always fixes nested GPU device maps and makes them into a single gpu device map that is easier to reason about, so maybe we merge that PR, and then ensure there are no nested gpu device maps when we run this?

Also if we have > 3 dimensions then then all upper dimensions are flattened into a z index, and then this loop should be ok.

[ThrudPrimrose]

Although I do not see an issue in the implementation (and well let's be honest, if there is a bug here it will come out really fast and we can then fix it), I think the organization of the files is very hard for me.

I would propose something like this, especially because we expect to have more access policies implemented over time:

dace/
├── sdfg/
│   ├── memlet_access_policy.py          # MemletAccessPolicy base class
│   └── memlet_access_policies/
│       ├── __init__.py
│       ├── copy_on_access.py            # CopyOnAccess
│       └── loop_cursor.py                   # LoopCursor
│
└── transformation/
    └── passes/
        ├── memlet_schedules/
        │   ├── __init__.py
        │   ├── helpers.py                          # Global memlet-path utilities
        │   └── assign_loop_cursors.py   # AssignLoopCursors
        │
        └── memlet_access_policy_lowering.py # LowerMemletAccessPolicies

[ThrudPrimrose]

As a note current structure is like this:

dace/
├── memlet.py
│   └── Memlet.access_policy
│
├── sdfg/
│   └── memlet_access_policy.py (can be split)
│       ├── MemletAccessPolicy
│       ├── CopyOnAccess
│       └── LoopCursor
│
│
└── transformation/
    └── passes/
        ├── memlet_access_policies.py (Imho feels too long)
        │   ├── AssignLoopCursors
        │   ├── LowerMemletAccessPolicies
        │   └── Loop-cursor lowering helpers
        │
        │
        └── __init__.py
            └── Exports both access-policy passes