[llvm] [AMDGPU] Guard RewriteMFMAFormStage recolor against unsafe def/use (PR #217396)

[Romanov Vlad via llvm-commits]

================
@@ -0,0 +1,87 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=amdgpu9.50-amd-amdhsa -run-pass=machine-scheduler \
+# RUN:     -amdgpu-disable-rewrite-mfma-form-sched-stage=false \
+# RUN:     -verify-machineinstrs -o - %s | FileCheck %s
+
+# Correctness test for the RewriteMFMAFormStage def-side recolor guard. The loop
+# MFMA's src2 %acc is partially redefined in the epilogue (%acc.sub0) and read by
+# a VGPR-only DS_READ_B32. isRecolorSafe scans the reaching uses of %acc's defs,
+# sees that read, and keeps %acc in VGPR -- bridging it to AGPR with a copy for
+# src2 instead of recoloring the whole vreg (which would abort with
+# "AReg_64_Align2.sub0 cannot be used for VGPR_32 operands").
+
+--- |
+  define amdgpu_kernel void @mfma_fix_recolor() #0 {
+    ret void
+  }
+
+  attributes #0 = { "amdgpu-flat-work-group-size"="1,256" }
+...
+---
+name:            mfma_fix_recolor
+tracksRegLiveness: true
+body: |
+  ; CHECK-LABEL: name: mfma_fix_recolor
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $vgpr0_vgpr1, $vgpr2_vgpr3, $sgpr0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   %s0:vreg_64_align2 = COPY $vgpr0_vgpr1
+  ; CHECK-NEXT:   %s1:vreg_64_align2 = COPY $vgpr2_vgpr3
+  ; CHECK-NEXT:   %cnt:sgpr_32 = COPY $sgpr0
+  ; CHECK-NEXT:   %acc:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:areg_64_align2 = COPY %acc
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   dead %mdst:areg_64_align2 = V_MFMA_F64_4X4X4F64_e64 %s0, %s1, [[COPY]], 0, 0, 0, implicit $mode, implicit $exec
+  ; CHECK-NEXT:   %cnt:sgpr_32 = S_SUB_I32 %cnt, 1, implicit-def $scc
+  ; CHECK-NEXT:   S_CBRANCH_SCC1 %bb.1, implicit $scc
+  ; CHECK-NEXT:   S_BRANCH %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   %acc.sub0:vreg_64_align2 = AV_MOV_B32_IMM_PSEUDO 7, implicit $exec
+  ; CHECK-NEXT:   dead %conn:vreg_64_align2 = COPY %acc
+  ; CHECK-NEXT:   %ld:vgpr_32 = DS_READ_B32_gfx9 %acc.sub0, 0, 0, implicit $exec :: (load (s32), addrspace 3)
+  ; CHECK-NEXT:   [[DEF:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF1:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF2:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF3:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF4:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF5:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF6:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF7:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   [[DEF8:%[0-9]+]]:vreg_1024 = IMPLICIT_DEF
+  ; CHECK-NEXT:   S_ENDPGM 0, implicit %ld, implicit [[DEF]], implicit [[DEF1]], implicit [[DEF2]], implicit [[DEF3]], implicit [[DEF4]], implicit [[DEF5]], implicit [[DEF6]], implicit [[DEF7]], implicit [[DEF8]]
+  bb.0:
+    successors: %bb.1
+    liveins: $vgpr0_vgpr1, $vgpr2_vgpr3, $sgpr0
+
+    %s0:vreg_64_align2 = COPY $vgpr0_vgpr1
+    %s1:vreg_64_align2 = COPY $vgpr2_vgpr3
+    %cnt:sgpr_32 = COPY $sgpr0
+    %p0:vreg_1024 = IMPLICIT_DEF
+    %p1:vreg_1024 = IMPLICIT_DEF
+    %p2:vreg_1024 = IMPLICIT_DEF
+    %p3:vreg_1024 = IMPLICIT_DEF
+    %p4:vreg_1024 = IMPLICIT_DEF
+    %p5:vreg_1024 = IMPLICIT_DEF
+    %p6:vreg_1024 = IMPLICIT_DEF
+    %p7:vreg_1024 = IMPLICIT_DEF
+    %p8:vreg_1024 = IMPLICIT_DEF
+    %acc:vreg_64_align2 = AV_MOV_B64_IMM_PSEUDO 0, implicit $exec
+
+  bb.1:
+    successors: %bb.1, %bb.2
+    dead %mdst:vreg_64_align2 = V_MFMA_F64_4X4X4F64_vgprcd_e64 %s0, %s1, %acc, 0, 0, 0, implicit $mode, implicit $exec
+    %cnt:sgpr_32 = S_SUB_I32 %cnt, 1, implicit-def $scc
+    S_CBRANCH_SCC1 %bb.1, implicit $scc
+    S_BRANCH %bb.2
+
+  bb.2:
+    %acc.sub0 = AV_MOV_B32_IMM_PSEUDO 7, implicit $exec
+    %ld:vgpr_32 = DS_READ_B32_gfx9 %acc.sub0, 0, 0, implicit $exec :: (load (s32), addrspace 3)
+    dead %conn:vreg_64_align2 = COPY %acc
----------------
romanovvlad wrote:

Minor. It seems we do not need this copy here.

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