[Codegen][GPU] Add TargetDetails for Nvidia Ada Arch (#24616)
Ada architecture seems to be slightly different to the previous
generation. With a different number of streaming multiprocessors. This
PR adds support for RTX 4000s series GPUs.
diff --git a/compiler/plugins/target/CUDA/test/target_device_features.mlir b/compiler/plugins/target/CUDA/test/target_device_features.mlir
index e9683f5..91d4f40 100644
--- a/compiler/plugins/target/CUDA/test/target_device_features.mlir
+++ b/compiler/plugins/target/CUDA/test/target_device_features.mlir
@@ -3,7 +3,7 @@
// RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=cuda},iree-hal-transformation-pipeline{serialize-executables=false})' \
// RUN: --iree-cuda-target=ada %s | FileCheck %s --check-prefix=SM89
// RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=cuda},iree-hal-transformation-pipeline{serialize-executables=false})' \
-// RUN: --iree-cuda-target=rtx4090 %s | FileCheck %s --check-prefix=SM89
+// RUN: --iree-cuda-target=rtx4090 %s | FileCheck %s --check-prefixes=SM89,RTX4090
// RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=cuda},iree-hal-transformation-pipeline{serialize-executables=false})' \
// RUN: --iree-cuda-target=sm_89 --iree-cuda-target-features=+ptx80 %s | FileCheck %s --check-prefix=PTX80
// RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=cuda},iree-hal-transformation-pipeline{serialize-executables=false})' \
@@ -16,8 +16,9 @@
// SM89-SAME: subgroup = shuffle|arithmetic, dot = dp4xi8toi32,
// SM89-SAME: mma = [<NV_MMA_SYNC_F32_16x8x16_F16>, <NV_MMA_SYNC_F16_16x8x16_F16>, <NV_MMA_SYNC_F32_16x8x16_BF16>, <NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>],
// SM89-SAME: subgroup_size_choices = [32], max_workgroup_sizes = [1024, 1024, 1024],
-// SM89-SAME: max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 166912,
-// SM89-SAME: max_workgroup_counts = [2147483647, 65535, 65535]>>
+// SM89-SAME: max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 101376,
+// SM89-SAME: max_workgroup_counts = [2147483647, 65535, 65535]>
+// RTX4090: chip = <wgp_count = 128>
// PTX80: target_info = #iree_gpu.target<arch = "sm_89", features = "+ptx80",
// SM120: target_info = #iree_gpu.target<arch = "sm_120", features = "+ptx87",
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
index 680d248..64ccd85 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
@@ -1019,6 +1019,35 @@
return &sm121Wgp;
}
+const WgpDetails *getAdaWgpDetails() {
+ // Ada (sm_89) exposes the same tensor-core instructions as Ampere-class
+ // hardware.
+ static const MMAIntrinsic mmaOps[] = {
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_F16,
+ MMAIntrinsic::NV_MMA_SYNC_F16_16x8x16_F16,
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16,
+ MMAIntrinsic::NV_WMMA_F32_16x16x16_F16,
+ MMAIntrinsic::NV_WMMA_F16_16x16x16_F16,
+ };
+ static const WgpDetails adaWgp = {allComputeBits,
+ allStorageBits,
+ allSubgroupOps,
+ allDotProductOps,
+ std::size(mmaOps),
+ mmaOps,
+ 0,
+ nullptr,
+ {32, 32},
+ {1024, 1024, 1024},
+ 1024,
+ // sm_89 caps shared memory at 99 KB per
+ // block (100 KB per SM); the 128 KB figure
+ // is the unified L1+shared size.
+ 99 * 1024,
+ {0x7fffffff, 0xffff, 0xffff}};
+ return &adaWgp;
+}
+
// Reports Ampere-class NVIDIA tensor core capabilities for GPU target
// selection.
const WgpDetails *getAmpereWgpDetails() {
@@ -1136,12 +1165,17 @@
std::optional<TargetDetails> getNVIDIAGPUTargetDetails(StringRef target) {
const WgpDetails *sm120Wgp = getSM120WgpDetails();
const WgpDetails *sm121Wgp = getSM121WgpDetails();
+ const WgpDetails *adaWgp = getAdaWgpDetails();
const WgpDetails *ampereWgp = getAmpereWgpDetails();
const WgpDetails *turingWgp = getTuringWgpDetails();
const WgpDetails *voltaWgp = getVoltaWgpDetails();
const WgpDetails *pascalWgp = getPascalWgpDetails();
static const ChipDetails a100Chip = {108};
+ static const ChipDetails rtx4090Chip = {128};
+ static const ChipDetails rtx4080Chip = {76};
+ static const ChipDetails rtx4070Chip = {46};
+ static const ChipDetails rtx4060Chip = {24};
static const ChipDetails rtx3090tiChip = {84};
static const ChipDetails rtx3090Chip = {82};
static const ChipDetails rtx3080tiChip = {80};
@@ -1153,15 +1187,18 @@
// lists mappings from microarchitectures to compute capabilities.
std::string lowerTarget = target.lower();
- // Treat RTX 40 family aliases as Ada-class devices when no exact chip is
- // named.
- if (StringRef(lowerTarget).starts_with("rtx40")) {
- return TargetDetails{ampereWgp, nullptr};
- }
return llvm::StringSwitch<std::optional<TargetDetails>>(lowerTarget)
// https://www.techpowerup.com/gpu-specs/a100-sxm4-80-gb.c3746
.Case("a100", TargetDetails{ampereWgp, &a100Chip})
+ // https://www.techpowerup.com/gpu-specs/geforce-rtx-4090.c3889
+ .Case("rtx4090", TargetDetails{adaWgp, &rtx4090Chip})
+ // https://www.techpowerup.com/gpu-specs/geforce-rtx-4080.c3888
+ .Case("rtx4080", TargetDetails{adaWgp, &rtx4080Chip})
+ // https://www.techpowerup.com/gpu-specs/geforce-rtx-4070.c3924
+ .Case("rtx4070", TargetDetails{adaWgp, &rtx4070Chip})
+ // https://www.techpowerup.com/gpu-specs/geforce-rtx-4060.c4107
+ .Case("rtx4060", TargetDetails{adaWgp, &rtx4060Chip})
// https://www.techpowerup.com/gpu-specs/geforce-rtx-3090-ti.c3829
.Case("rtx3090ti", TargetDetails{ampereWgp, &rtx3090tiChip})
// https://www.techpowerup.com/gpu-specs/geforce-rtx-3090.c3622
@@ -1179,7 +1216,7 @@
// validation.
.Case("sm_120", TargetDetails{sm120Wgp, nullptr})
.Case("sm_121", TargetDetails{sm121Wgp, nullptr})
- .Cases({"ada", "sm_89"}, TargetDetails{ampereWgp, nullptr})
+ .Cases({"ada", "sm_89"}, TargetDetails{adaWgp, nullptr})
.Cases({"ampere", "sm_80", "sm_86", "sm_87"},
TargetDetails{ampereWgp, nullptr})
.Cases({"turing", "sm_75"}, TargetDetails{turingWgp, nullptr})
