Skip to content

Commit

Permalink
[stdlib][GPU] Fix invalid constant addr space
Browse files Browse the repository at this point in the history
On Nvidia, address space value of 2 in the IR is reserved.  So, it's
invalid for us to represent the constant address space with the value of
`2`.

Good news: nobody is setting this value currently.  So, just remove it.
While here, rename the existing `PARAM` address space name to be called
"CONSTANT" instead, which is more consistent with the
[docs](https://docs.nvidia.com/cuda/nvvm-ir-spec/#address-space).  Note
that AMD calls address space 4 "constant" as well.

MODULAR_ORIG_COMMIT_REV_ID: 35bbc8ad2b82efdddf6fa4cf537132af2ca2d3af
  • Loading branch information
JoeLoser authored and modularbot committed Dec 11, 2024
1 parent ef4e86f commit 3efbf09
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 6 deletions.
7 changes: 5 additions & 2 deletions stdlib/src/memory/memory.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -414,9 +414,12 @@ fn stack_allocation[

@parameter
if is_gpu():
# On NVGPU, SHARED and PARAM address spaces lower to global memory.
# On NVGPU, SHARED and CONSTANT address spaces lower to global memory.
@parameter
if address_space in (_GPUAddressSpace.SHARED, _GPUAddressSpace.PARAM):
if address_space in (
_GPUAddressSpace.SHARED,
_GPUAddressSpace.CONSTANT,
):
alias global_name = name.value() if name else "_global_alloc"
return __mlir_op.`pop.global_alloc`[
name = global_name.value,
Expand Down
6 changes: 2 additions & 4 deletions stdlib/src/memory/pointer.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,10 @@ struct _GPUAddressSpace(EqualityComparable):
"""Generic address space."""
alias GLOBAL = AddressSpace(1)
"""Global address space."""
alias CONSTANT = AddressSpace(2) if is_nvidia_gpu() else AddressSpace(4)
"""Constant address space."""
alias SHARED = AddressSpace(3)
"""Shared address space."""
alias PARAM = AddressSpace(4)
"""Param address space."""
alias CONSTANT = AddressSpace(4)
"""Constant address space."""
alias LOCAL = AddressSpace(5) if is_nvidia_gpu() else AddressSpace(3)
"""Local address space."""

Expand Down

0 comments on commit 3efbf09

Please sign in to comment.