[llvm] [BasicAA] Treat different VScale intrinsics as the same value. (PR #81152)

[David Green via llvm-commits]

https://github.com/davemgreen updated https://github.com/llvm/llvm-project/pull/81152

>From b05d70eb75e1b69f4416473ae0a56e6146ba6ddc Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Mon, 12 Feb 2024 10:22:32 +0000
Subject: [PATCH] [BasicAA] Treat different VScale intrinsics as the same
 value.

The IR may contain multiple llvm.vscale intrinsics that have not been CSEd.
This patch ensures that multiple vscales can be treated the same, either in the
decomposition of geps and when we subtrack one decomposition from another.
---
 llvm/lib/Analysis/BasicAliasAnalysis.cpp | 12 ++++++++++--
 llvm/test/Analysis/BasicAA/vscale.ll     |  4 ++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Analysis/BasicAliasAnalysis.cpp b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
index 2ea29177a5cf0f..8dfc4b2a1cb1f4 100644
--- a/llvm/lib/Analysis/BasicAliasAnalysis.cpp
+++ b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
@@ -188,6 +188,12 @@ static bool isObjectSize(const Value *V, TypeSize Size, const DataLayout &DL,
   return ObjectSize && *ObjectSize == Size;
 }
 
+/// Return true if both V1 and V2 are VScale
+static bool areBothVScale(const Value *V1, const Value *V2) {
+  return PatternMatch::match(V1, PatternMatch::m_VScale()) &&
+         PatternMatch::match(V2, PatternMatch::m_VScale());
+}
+
 //===----------------------------------------------------------------------===//
 // CaptureInfo implementations
 //===----------------------------------------------------------------------===//
@@ -679,7 +685,8 @@ BasicAAResult::DecomposeGEPExpression(const Value *V, const DataLayout &DL,
       //   A[x][x] -> x*16 + x*4 -> x*20
       // This also ensures that 'x' only appears in the index list once.
       for (unsigned i = 0, e = Decomposed.VarIndices.size(); i != e; ++i) {
-        if (Decomposed.VarIndices[i].Val.V == LE.Val.V &&
+        if ((Decomposed.VarIndices[i].Val.V == LE.Val.V ||
+             areBothVScale(Decomposed.VarIndices[i].Val.V, LE.Val.V)) &&
             Decomposed.VarIndices[i].Val.hasSameCastsAs(LE.Val)) {
           Scale += Decomposed.VarIndices[i].Scale;
           LE.IsNSW = false; // We cannot guarantee nsw for the merge.
@@ -1792,7 +1799,8 @@ void BasicAAResult::subtractDecomposedGEPs(DecomposedGEP &DestGEP,
     bool Found = false;
     for (auto I : enumerate(DestGEP.VarIndices)) {
       VariableGEPIndex &Dest = I.value();
-      if (!isValueEqualInPotentialCycles(Dest.Val.V, Src.Val.V, AAQI) ||
+      if ((!isValueEqualInPotentialCycles(Dest.Val.V, Src.Val.V, AAQI) &&
+           !areBothVScale(Dest.Val.V, Src.Val.V)) ||
           !Dest.Val.hasSameCastsAs(Src.Val))
         continue;
 
diff --git a/llvm/test/Analysis/BasicAA/vscale.ll b/llvm/test/Analysis/BasicAA/vscale.ll
index 18aad4ffe73f49..05b9b6b3c3a970 100644
--- a/llvm/test/Analysis/BasicAA/vscale.ll
+++ b/llvm/test/Analysis/BasicAA/vscale.ll
@@ -488,9 +488,9 @@ define void @onevscale(ptr %p) vscale_range(1,16) {
 }
 
 ; CHECK-LABEL: twovscales
-; CHECK-DAG:   MayAlias:     <vscale x 4 x i32>* %vp161, <vscale x 4 x i32>* %vp162
+; CHECK-DAG:   MustAlias:    <vscale x 4 x i32>* %vp161, <vscale x 4 x i32>* %vp162
 ; CHECK-DAG:   NoAlias:      <vscale x 4 x i32>* %vp161, <vscale x 4 x i32>* %vp161b
-; CHECK-DAG:   MayAlias:     <vscale x 4 x i32>* %vp161b, <vscale x 4 x i32>* %vp162
+; CHECK-DAG:   NoAlias:      <vscale x 4 x i32>* %vp161b, <vscale x 4 x i32>* %vp162
 define void @twovscales(ptr %p) vscale_range(1,16) {
   %v1 = call i64 @llvm.vscale.i64()
   %v2 = call i64 @llvm.vscale.i64()