refactor[cartesian]: unexpanded sdfg cleanups

[romanc]

Description

Refactors from recent debugging sessions around transient arrays in the "unexpanded SDFG" (the one with the library nodes):

• Remove unused **kwargs from OirSDFGBuilder
• Forward debugging information about transient arrays to DaCe
• Use a (constant) variable for connector prefixes of data going in to/out of the library nodes
• Configure the lifetime of transient arrays directly in OirSDFGBuilder
• Configure storage type of transient arrays directly in OirSDFGBuilder (for GPU targets)
• Configure the LibraryNode's device type directly in OirSDFGBuilder (for GPU targets)
• Enum type for StorageDevice in LayoutInfo

Sidenote on the allocation lifetime: In the orchestrated code path, we reset the allocation lifetime of transients to SDFG when we freeze the stencil with origin/domain

gt4py/src/gt4py/cartesian/backend/dace_backend.py

Lines 270 to 317 in ac253b6

	def freeze_origin_domain_sdfg(inner_sdfg, arg_names, field_info, *, origin, domain):
	wrapper_sdfg = dace.SDFG("frozen_" + inner_sdfg.name)
	state = wrapper_sdfg.add_state("frozen_" + inner_sdfg.name + "_state")
	inputs = set()
	outputs = set()
	for inner_state in inner_sdfg.nodes():
	for node in inner_state.nodes():
	if (
	not isinstance(node, dace.nodes.AccessNode)
	or inner_sdfg.arrays[node.data].transient
	):
	continue
	if node.has_reads(inner_state):
	inputs.add(node.data)
	if node.has_writes(inner_state):
	outputs.add(node.data)
	nsdfg = state.add_nested_sdfg(inner_sdfg, None, inputs, outputs)
	_sdfg_add_arrays_and_edges(
	field_info, wrapper_sdfg, state, inner_sdfg, nsdfg, inputs, outputs, origins=origin
	)
	# in special case of empty domain, remove entire SDFG.
	if any(d == 0 for d in domain):
	states = wrapper_sdfg.states()
	assert len(states) == 1
	for node in states[0].nodes():
	state.remove_node(node)
	# make sure that symbols are passed through to inner sdfg
	for symbol in nsdfg.sdfg.free_symbols:
	if symbol not in wrapper_sdfg.symbols:
	wrapper_sdfg.add_symbol(symbol, nsdfg.sdfg.symbols[symbol])
	# Try to inline wrapped SDFG before symbols are specialized to avoid extra views
	inline_sdfgs(wrapper_sdfg)
	_sdfg_specialize_symbols(wrapper_sdfg, domain)
	for _, _, array in wrapper_sdfg.arrays_recursive():
	if array.transient:
	array.lifetime = dace.dtypes.AllocationLifetime.SDFG
	wrapper_sdfg.arg_names = arg_names
	return wrapper_sdfg

This might be relevant when tracking down orchestration performance. Seems odd at least.

Requirements

• All fixes and/or new features come with corresponding tests.
  Covered by existing tests
• Important design decisions have been documented in the appropriate ADR inside the docs/development/ADRs/ folder.
  N/A

[romanc]

(rebased onto main branch because there was a hickup on the CSCS-CI)

[FlorianDeconinck]

Re: freeze_origin_domain_sdfg

My cursory look seems to point that this is only called at top level for DaceLazyStencil et al. e.g. it's not called when we are doing orchestration but for the top level arguments (and then it's fine).

We should log a task to double check and we need a methodological/code way to differentiate code path we expect is in stencil mode and what is in orchestrated mode

[FlorianDeconinck]

Looking good. Couple of question.

Do the enum for the device type, the string is making me sad.

[FlorianDeconinck]

LGTM

[FlorianDeconinck]

@havogt Care to review or ok to push?

[havogt]

I would prefer to use the DeviceType enum from _core.definitions instead of introducing a new enum. Can you check if that works?

[havogt]

It would be nice to use the DeviceType from _core.definitions. I don't have time to check how ROCM would fit here.
Maybe we can do

CurrentGPU = core_defs.DeviceType.ROCM if cp.cuda.runtime.is_hip else core_defs.DeviceType.CUDA

somewhere. cupy currently has a limitation that it is either installed for cuda or for rocm, so it would be fine.

[romanc]

Sure, let's reuse DeviceType from the core definitions. I honestly don't know what our support level for AMD graphic cards is. From what I read in the code, it looks like that's gonna be an item on the tasklist at some point. So I personally wouldn't worry too much about ROCM for now

[FlorianDeconinck]

ECMWF runs AMD hardware - so we need to keep this alive

[romanc]

For me to understand, which backend is compatible with AMD hardware? From what I read, the CudaBackend as well as the DaceGPUBackend, generate cuda code.

[FlorianDeconinck]

So DaCe can HIP - I believe that's where they are getting there AMD goodies. @stubbiali, can you talk to the AMD hardware use? Am I misremembering ?

[stubbiali]

Both GridTools and DaCe GPU backends support AMD graphics cards. We have been running GT4Py codes on LUMI for at least 1.5 years.

[romanc]

Putting this back to draft mode.

• Re-using DeviceType from _core.definitions is possible and I think it's worth doing.
• It's a bit of work and I might need to re-shuffle some things because tach complains that gt4py.cartesian should not depend on gt4py._core directly. I guess everything related to cpu/gpu device(s) should go through the storage module. In the end, there are many checks like "if this is a gpu device, then do this ..." and they should be centralized (in the gt_storage package or wherever) such that we can deal with CUDA / ROCM in a single place.
• Mixing the original changes in this PR with changing the "cpu" / "gpu" strings re-use DeviceType is going to look ugly in the git history (especially when squashed).
• I'll push a WIP state here and then probably separate the two things into separate (consecutive) PRs tomorrow.

[romanc]

Closing as superseded by #1860 and #1867.