From 3f713f5f4743ddb6715e4ea4a361784b54489e5a Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub@nod-labs.com>
Date: Wed, 29 Jan 2025 12:49:42 -0500
Subject: [PATCH] [ROCm] Add mi325x to known targets (#19846)

---
 .../plugins/target/ROCM/test/target_device_features.mlir     | 5 ++++-
 .../Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp         | 4 +++-
 .../docs/guides/deployment-configurations/gpu-rocm.md        | 1 +
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/compiler/plugins/target/ROCM/test/target_device_features.mlir b/compiler/plugins/target/ROCM/test/target_device_features.mlir
index c50bf83e6330..595b6270aac1 100644
--- a/compiler/plugins/target/ROCM/test/target_device_features.mlir
+++ b/compiler/plugins/target/ROCM/test/target_device_features.mlir
@@ -4,6 +4,8 @@
 // RUN:   --iree-hip-target=mi300a %s | FileCheck %s --check-prefixes=GFX942,MI300A
 // RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=hip},iree-hal-transformation-pipeline{serialize-executables=false})' \
 // RUN:   --iree-hip-target=mi308x %s | FileCheck %s --check-prefixes=GFX942,MI308X
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=hip},iree-hal-transformation-pipeline{serialize-executables=false})' \
+// RUN:   --iree-hip-target=mi325x %s | FileCheck %s --check-prefixes=GFX942,MI325X
 //
 // RUN: iree-opt --pass-pipeline='builtin.module(iree-hal-assign-target-devices{targetDevices=hip},iree-hal-transformation-pipeline{serialize-executables=false})' \
 // RUN:   --iree-hip-target=gfx941 --iree-hip-target-features=+sramecc,-xnack %s | FileCheck %s --check-prefix=GFX941
@@ -32,7 +34,8 @@
 // GFX942-SAME:         max_workgroup_counts = [2147483647, 2147483647, 2147483647],
 // MI300X: chip = <wgp_count = 304, sku = "mi300x">>
 // MI300A: chip = <wgp_count = 228, sku = "mi300a">>
-// MI308X: chip = <wgp_count = 80,  sku = "mi308x">>
+// MI308X: chip = <wgp_count = 80, sku = "mi308x">>
+// MI325X: chip = <wgp_count = 304, sku = "mi325x">>
 
 // GFX941: target = #iree_gpu.target<arch = "gfx941",
 // GFX941-SAME:         features = "+sramecc,-xnack"
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 f99f4368d895..8e6e8f949907 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
@@ -287,6 +287,7 @@ std::optional<TargetDetails> getAMDGPUTargetDetails(StringRef target) {
   static const ChipDetails mi300xChip = {304, "mi300x"};
   static const ChipDetails mi300aChip = {228, "mi300a"};
   static const ChipDetails mi308xChip = {80, "mi308x"};
+  static const ChipDetails mi325xChip = {304, "mi325x"};
 
   // "AMD Instinct MI200 Series Accelerator Product Offerings" in Page 14 of
   // https://www.amd.com/content/dam/amd/en/documents/instinct-business-docs/white-papers/amd-cdna2-white-paper.pdf
@@ -310,6 +311,7 @@ std::optional<TargetDetails> getAMDGPUTargetDetails(StringRef target) {
   // See https://llvm.org/docs/AMDGPUUsage.html#processors for gfxN to
   // cdnaN/rdnaN mapping.
   return llvm::StringSwitch<std::optional<TargetDetails>>(target.lower())
+      .Case("mi325x", TargetDetails{cdna3Wgp, &mi325xChip})
       .Case("mi300x", TargetDetails{cdna3Wgp, &mi300xChip})
       .Case("mi300a", TargetDetails{cdna3Wgp, &mi300aChip})
       .Case("mi308x", TargetDetails{cdna3Wgp, &mi308xChip})
@@ -356,7 +358,7 @@ StringRef normalizeAMDGPUTarget(StringRef target) {
     return target;
 
   return llvm::StringSwitch<StringRef>(target.lower())
-      .Cases("mi300a", "mi300x", "mi308x", "gfx942")
+      .Cases("mi300a", "mi300x", "mi308x", "mi325x", "gfx942")
       .Cases("mi250x", "mi250", "mi210", "cdna2", "gfx90a")
       .Cases("mi100", "cdna1", "gfx908")
       .Cases("rx7900xtx", "rx7900xt", "w7900", "w7800", "gfx1100")
diff --git a/docs/website/docs/guides/deployment-configurations/gpu-rocm.md b/docs/website/docs/guides/deployment-configurations/gpu-rocm.md
index f8c499b6bd36..20d34388ecf1 100644
--- a/docs/website/docs/guides/deployment-configurations/gpu-rocm.md
+++ b/docs/website/docs/guides/deployment-configurations/gpu-rocm.md
@@ -128,6 +128,7 @@ compile towards each GPU chip. Here is a table of commonly used architectures:
 | AMD MI300A               | `mi300a`    | `gfx942`            | `cdna3`                |
 | AMD MI300X               | `mi300x`    | `gfx942`            | `cdna3`                |
 | AMD MI308X               | `mi308x`    | `gfx942`            | `cdna3`                |
+| AMD MI325X               | `mi325x`    | `gfx942`            | `cdna3`                |
 | AMD RX7900XTX            | `rx7900xtx` | `gfx1100`           | `rdna3`                |
 | AMD RX7900XT             | `rx7900xt`  | `gfx1100`           | `rdna3`                |
 | AMD PRO W7900            | `w7900`     | `gfx1100`           | `rdna3`                |