Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support TMA with 64-bit indexing #3599

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Support TMA with 64-bit indexing #3599

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2e66b80
Support TMA with 64-bit indexing
jacobhinkle Dec 16, 2024
ade644f
Merge remote-tracking branch 'origin/main' into tma_int64_indexing
jacobhinkle Dec 17, 2024
5ec3ee5
Add test
jacobhinkle Dec 17, 2024
308d4fe
Cast to int32_t instead of uint32_t
jacobhinkle Dec 17, 2024
1df92f9
Add check for TMA axes when using 64-bit indexing
jacobhinkle Dec 17, 2024
9bdf4c8
Fix type of literal that caused overflow in check
jacobhinkle Dec 17, 2024
d1c71c6
Cast in the IR instead of the runtime memory.cu
jacobhinkle Dec 17, 2024
8e55b54
Fix inequality in check
jacobhinkle Dec 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions csrc/runtime/executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -375,12 +375,6 @@ void KernelExecutor::compile(
!(compile_params.index_type.value() == PrimDataType::Int32 &&
arg_index_type == PrimDataType::Int),
"Compilation with int32 is requested but int64 is required for the arguments");
NVF_ERROR(
!has_cp_async_bulk ||
(compile_params.index_type.value() == PrimDataType::Int32),
"Compilation with int64 is requested but int32 is required because ",
"of TMA operations.");

} else if (arg_index_type == PrimDataType::Int) {
// If the given compile option doesn't specify the index type, and
// the arguments require 64-bit indexing, we need to use 64-bit
Expand Down
35 changes: 19 additions & 16 deletions runtime/memory.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ __device__ inline bool electSync(const uint32_t& membermask) {
template <int dim>
struct CpAsyncBulkTensorTileG2SIndex {
const TensorMap* descriptor;
Array<int32_t, dim> crds;
Array<nvfuser_index_t, dim> crds;
uint32_t mbarrier;
};

Expand All @@ -108,7 +108,10 @@ __device__ inline void cpAsyncBulkTensorTileG2S(
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3}], [%2];"
:
: "r"(smem_addr), "l"(gmem_int_desc), "r"(src.mbarrier), "r"(src.crds[0])
: "r"(smem_addr),
"l"(gmem_int_desc),
"r"(src.mbarrier),
"r"((uint32_t)src.crds[0])
: "memory");
}

Expand All @@ -123,8 +126,8 @@ __device__ inline void cpAsyncBulkTensorTileG2S(
: "r"(smem_addr),
"l"(gmem_int_desc),
"r"(src.mbarrier),
"r"(src.crds[0]),
"r"(src.crds[1])
"r"((uint32_t)src.crds[0]),
"r"((uint32_t)src.crds[1])
: "memory");
}

Expand All @@ -139,9 +142,9 @@ __device__ inline void cpAsyncBulkTensorTileG2S(
: "r"(smem_addr),
"l"(gmem_int_desc),
"r"(src.mbarrier),
"r"(src.crds[0]),
"r"(src.crds[1]),
"r"(src.crds[2])
"r"((uint32_t)src.crds[0]),
"r"((uint32_t)src.crds[1]),
"r"((uint32_t)src.crds[2])
: "memory");
}

Expand All @@ -156,10 +159,10 @@ __device__ inline void cpAsyncBulkTensorTileG2S(
: "r"(smem_addr),
"l"(gmem_int_desc),
"r"(src.mbarrier),
"r"(src.crds[0]),
"r"(src.crds[1]),
"r"(src.crds[2]),
"r"(src.crds[3])
"r"((uint32_t)src.crds[0]),
"r"((uint32_t)src.crds[1]),
"r"((uint32_t)src.crds[2]),
"r"((uint32_t)src.crds[3])
: "memory");
}

Expand All @@ -174,11 +177,11 @@ __device__ inline void cpAsyncBulkTensorTileG2S(
: "r"(smem_addr),
"l"(gmem_int_desc),
"r"(src.mbarrier),
"r"(src.crds[0]),
"r"(src.crds[1]),
"r"(src.crds[2]),
"r"(src.crds[3]),
"r"(src.crds[4])
"r"((uint32_t)src.crds[0]),
"r"((uint32_t)src.crds[1]),
"r"((uint32_t)src.crds[2]),
"r"((uint32_t)src.crds[3]),
"r"((uint32_t)src.crds[4])
: "memory");
}

Expand Down
Loading