[llvm] [Verifier] Verify AMX tile-register index operands are in range (PR #204890)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-ir

Author: jofrn

<details>
<summary>Changes</summary>

AMX has 8 physical tile registers (TMM0-TMM7), so the tile-index operands
of the AMX intrinsics must be in [0, 8): operand 0 for the tile
load/store/zero intrinsics, operands 0-2 for the tdp* family.

---
Full diff: https://github.com/llvm/llvm-project/pull/204890.diff


2 Files Affected:

- (modified) llvm/lib/IR/Verifier.cpp (+24) 
- (added) llvm/test/Verifier/x86-amx-tile-register-index.ll (+30) 


``````````diff
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 48cfec0f55c41..2754fe49eef0a 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -99,6 +99,7 @@
 #include "llvm/IR/IntrinsicsARM.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/IntrinsicsWebAssembly.h"
+#include "llvm/IR/IntrinsicsX86.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MemoryModelRelaxationAnnotations.h"
 #include "llvm/IR/Metadata.h"
@@ -7095,6 +7096,29 @@ void Verifier::visitIntrinsicCall(Intrinsic::ID ID, CallBase &Call) {
           Call);
     break;
   }
+  case Intrinsic::x86_tilezero:
+  case Intrinsic::x86_tileloadd64:
+  case Intrinsic::x86_tileloaddt164:
+  case Intrinsic::x86_tilestored64: {
+    // AMX provides 8 physical tile registers, TMM0-TMM7; the leading immediate
+    // operand is a tile-register index and so must name one of them.
+    Check(cast<ConstantInt>(Call.getArgOperand(0))->getZExtValue() < 8,
+          "AMX tile register index must be in the range [0, 8)", Call);
+    break;
+  }
+  case Intrinsic::x86_tdpbssd:
+  case Intrinsic::x86_tdpbsud:
+  case Intrinsic::x86_tdpbusd:
+  case Intrinsic::x86_tdpbuud:
+  case Intrinsic::x86_tdpbf16ps:
+  case Intrinsic::x86_tdpfp16ps: {
+    // The destination and two source operands are each a tile-register index
+    // (TMM0-TMM7).
+    for (unsigned I : {0, 1, 2})
+      Check(cast<ConstantInt>(Call.getArgOperand(I))->getZExtValue() < 8,
+            "AMX tile register index must be in the range [0, 8)", Call);
+    break;
+  }
   case Intrinsic::callbr_landingpad: {
     const auto *CBR = dyn_cast<CallBrInst>(Call.getOperand(0));
     Check(CBR, "intrinstic requires callbr operand", &Call);
diff --git a/llvm/test/Verifier/x86-amx-tile-register-index.ll b/llvm/test/Verifier/x86-amx-tile-register-index.ll
new file mode 100644
index 0000000000000..47b17e25ad589
--- /dev/null
+++ b/llvm/test/Verifier/x86-amx-tile-register-index.ll
@@ -0,0 +1,30 @@
+; RUN: not opt -passes=verify -S < %s 2>&1 | FileCheck %s
+
+; AMX exposes 8 physical tile registers (TMM0-TMM7), so a tile-register index
+; operand must be in the range [0, 8). This is a structural constraint on the
+; immediate operand and needs no subtarget information.
+
+declare void @llvm.x86.tilezero(i8 immarg)
+declare void @llvm.x86.tileloadd64(i8 immarg, ptr, i64)
+declare void @llvm.x86.tdpbssd(i8 immarg, i8 immarg, i8 immarg)
+
+define void @bad_tilezero() {
+; CHECK: AMX tile register index must be in the range [0, 8)
+  call void @llvm.x86.tilezero(i8 8)
+  ret void
+}
+
+; Only one of the three tile operands is out of range.
+define void @bad_tdpbssd() {
+; CHECK: AMX tile register index must be in the range [0, 8)
+  call void @llvm.x86.tdpbssd(i8 0, i8 1, i8 9)
+  ret void
+}
+
+; A valid tile index (0) with a large *constant* stride. The stride is operand 2,
+; not a tile-register index, so it must not be diagnosed.
+; CHECK-NOT: AMX tile register index
+define void @good_tileloadd64(ptr %p) {
+  call void @llvm.x86.tileloadd64(i8 0, ptr %p, i64 64)
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/204890