Skip to content

Commit

Permalink
Fix the bad interval
Browse files Browse the repository at this point in the history
Remove not needed restriction on CUDA version
  • Loading branch information
FlorianDeconinck committed Jan 8, 2025
1 parent 8040178 commit 673fa3f
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 38 deletions.
15 changes: 0 additions & 15 deletions src/gt4py/cartesian/frontend/gtscript_frontend.py
Original file line number Diff line number Diff line change
Expand Up @@ -1451,21 +1451,6 @@ def visit_Assign(self, node: ast.Assign) -> list:
message="Assignment to non-zero offsets in K is not available in PARALLEL. Choose FORWARD or BACKWARD.",
loc=nodes.Location.from_ast_node(t),
)
if self.backend_name in ["gt:gpu", "dace:gpu"]:
import cupy as cp

if cp.cuda.runtime.runtimeGetVersion() < 12000:
raise GTScriptSyntaxError(
message=f"Assignment to non-zero offsets in K is not available in {self.backend_name} for CUDA<12. Please update CUDA.",
loc=nodes.Location.from_ast_node(t),
)

if self.backend_name in ["gt:gpu"]:
raise GTScriptSyntaxError(
message=f"Assignment to non-zero offsets in K is not available in {self.backend_name} as an unsolved bug remains."
"Please refer to https://github.com/GridTools/gt4py/issues/1754.",
loc=nodes.Location.from_ast_node(t),
)

if not self._is_known(name):
if name in self.temp_decls:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -582,17 +582,6 @@ def test_K_offset_write(backend):
# Cuda generates bad code for the K offset
if backend == "cuda":
pytest.skip("cuda K-offset write generates bad code")
if backend in ["dace:gpu"]:
import cupy as cp

if cp.cuda.runtime.runtimeGetVersion() < 12000:
pytest.skip(
f"{backend} backend with CUDA 11 and/or GCC 10.3 is not capable of K offset write, update CUDA/GCC if possible"
)
if backend in ["gt:gpu"]:
pytest.skip(
f"{backend} backend is not capable of K offset write, bug remains unsolved: https://github.com/GridTools/gt4py/issues/1754"
)

arraylib = get_array_library(backend)
array_shape = (1, 1, 4)
Expand Down Expand Up @@ -664,25 +653,14 @@ def backward(A: Field[np.float64], B: Field[np.float64], scalar: np.float64):
def test_K_offset_write_conditional(backend):
if backend == "cuda":
pytest.skip("Cuda backend is not capable of K offset write")
if backend in ["dace:gpu"]:
import cupy as cp

if cp.cuda.runtime.runtimeGetVersion() < 12000:
pytest.skip(
f"{backend} backend with CUDA 11 and/or GCC 10.3 is not capable of K offset write, update CUDA/GCC if possible"
)
if backend in ["gt:gpu"]:
pytest.skip(
f"{backend} backend is not capable of K offset write, bug remains unsolved: https://github.com/GridTools/gt4py/issues/1754"
)

arraylib = get_array_library(backend)
array_shape = (1, 1, 4)
K_values = arraylib.arange(start=40, stop=44)

@gtscript.stencil(backend=backend)
def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scalar: np.float64):
with computation(BACKWARD), interval(1, None):
with computation(BACKWARD), interval(1, -1):
if A > 0 and B > 0:
A[0, 0, -1] = scalar
B[0, 0, 1] = A
Expand All @@ -700,6 +678,42 @@ def column_physics_conditional(A: Field[np.float64], B: Field[np.float64], scala
backend=backend, aligned_index=(0, 0, 0), shape=array_shape, dtype=np.float64
)
column_physics_conditional(A, B, 2.0)
# Manual unroll of the above
# Starts with
# - A[...] = [40, 41, 42, 43]
# - B[...] = [1, 1, 1, 1]
# Now in-stencil
# ITERATION k = 2 of [2:1]
# if condition
# - A[2] == 42 && B[2] == 1 => True
# - A[1] = 2.0
# - B[3] = A[2] = 42
# while
# - lev = 1
# - A[2] == 42 && B[2] == 1 => True
# - A[3] = -1
# - B[2] = -1
# - lev = 2
# - A[2] == 42 && B[2] == -1 => False
# End of iteration state
# - A[...] = A[40, 2.0, 2.0, -1]
# - B[...] = A[1, 1, -1, 42]
# ITERATION k = 1 of [3:1]
# if condition
# - A[1] == 2.0 && B[1] == 1 => True
# - A[0] = 2.0
# - B[2] = A[1] = 2.0
# while
# - lev = 1
# - A[1] == 2.0 && B[1] == 1 => True
# - A[2] = -1
# - B[1] = -1
# - lev = 2
# - A[1] == 2.0 && B[2] == -1 => False
# End of stencil state
# - A[...] = A[2.0, 2.0, -1, -1]
# - B[...] = A[1, -1, 2.0, 42]

assert (A[0, 0, :] == arraylib.array([2, 2, -1, -1])).all()
assert (B[0, 0, :] == arraylib.array([1, -1, 2, 42])).all()

Expand Down

0 comments on commit 673fa3f

Please sign in to comment.