[llvm] [X86] Avoid assertion failure with implicit immediates for AMX tile dimensions (PR #174128)

[Thurston Dang via llvm-commits]

https://github.com/thurstond updated https://github.com/llvm/llvm-project/pull/174128

>From a36904e0af5ef72715db38587d27fe7ac11394ba Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Thu, 1 Jan 2026 05:02:41 +0000
Subject: [PATCH 1/2] [X86] Avoid assertion failure with implicit immediates
 for AMX tile dimensions

https://github.com/llvm/llvm-project/pull/165556 removed support for
implicit immediates in
llvm/include/llvm/CodeGen/TileShapeInfo.h::deduceImm(). Although
implicit immediates are generally not valid tile dimensions, they are
commonly encountered in reduced test cases, leading to assertion
failures.

This patch restores the support for implicit immediates (though does not
assume they will be zero).

Fixes: https://github.com/llvm/llvm-project/issues/174127
---
 llvm/include/llvm/CodeGen/TileShapeInfo.h | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/CodeGen/TileShapeInfo.h b/llvm/include/llvm/CodeGen/TileShapeInfo.h
index 24d9de842645a..96027a0bd5b2e 100644
--- a/llvm/include/llvm/CodeGen/TileShapeInfo.h
+++ b/llvm/include/llvm/CodeGen/TileShapeInfo.h
@@ -70,7 +70,14 @@ class ShapeT {
       for (const MachineOperand &DefMO : MRI->def_operands(Reg)) {
         const auto *MI = DefMO.getParent();
         if (MI->isMoveImmediate()) {
-          Imm = MI->getOperand(1).getImm();
+          if (MI->getOperand(1).isImm()) {
+            Imm = MI->getOperand(1).getImm();
+          } else {
+            assert(MI->getOperand(1).isImplicit() &&
+                   "Operand 1 is assumed to be implicit.");
+            // The implicit immediate can vary (MOV32r0, MOV32r1, MOV32r_1,
+            // ...) but in any case, is not a valid shape.
+          }
           break;
         }
       }

>From bd9fd9f5b23625e37e53985f96df93b51069bcfd Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Fri, 2 Jan 2026 19:34:20 +0000
Subject: [PATCH 2/2] Also revert change in
 llvm/lib/Target/X86/X86TileConfig.cpp

---
 llvm/lib/Target/X86/X86TileConfig.cpp | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/X86/X86TileConfig.cpp b/llvm/lib/Target/X86/X86TileConfig.cpp
index 09ef8fbc12de9..7702a70667ebd 100644
--- a/llvm/lib/Target/X86/X86TileConfig.cpp
+++ b/llvm/lib/Target/X86/X86TileConfig.cpp
@@ -170,7 +170,14 @@ bool X86TileConfig::runOnMachineFunction(MachineFunction &MF) {
                    "Cannot initialize with different shapes");
             continue;
           }
-          Imm = DefMI.getOperand(1).getImm();
+          if (DefMI.getOperand(1).isImm()) {
+            Imm = DefMI.getOperand(1).getImm();
+          } else {
+            assert(DefMI.getOpcode() == X86::MOV32r0 &&
+                   "The opcode is assumed to be MOV32r0 if the operand is not "
+                   "immediate.");
+            Imm = 0;
+          }
 
           NewMI = addFrameReference(
                       BuildMI(MF.front(), ++ConstMI->getIterator(), DL,