From a6fa58dda692a28847cd8bc3c48cae1057682001 Mon Sep 17 00:00:00 2001
From: Pietro Ghiglio <pietro.ghiglio@codeplay.com>
Date: Wed, 15 May 2024 15:03:21 +0200
Subject: [PATCH] [VPlan] Add scalar inferencing support for addrspace cast
 (#92107)

Fixes https://github.com/llvm/llvm-project/issues/91434

PR: https://github.com/llvm/llvm-project/pull/92107
---
 .../Transforms/Vectorize/VPlanAnalysis.cpp    |  1 +
 llvm/test/Transforms/LoopVectorize/as_cast.ll | 40 +++++++++++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 llvm/test/Transforms/LoopVectorize/as_cast.ll

diff --git a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
index 97a8a1803bbf5a5..2f9fb68469b9f44 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp
@@ -160,6 +160,7 @@ Type *VPTypeAnalysis::inferScalarTypeForRecipe(const VPReplicateRecipe *R) {
   case Instruction::ICmp:
   case Instruction::FCmp:
     return IntegerType::get(Ctx, 1);
+  case Instruction::AddrSpaceCast:
   case Instruction::Alloca:
   case Instruction::BitCast:
   case Instruction::Trunc:
diff --git a/llvm/test/Transforms/LoopVectorize/as_cast.ll b/llvm/test/Transforms/LoopVectorize/as_cast.ll
new file mode 100644
index 000000000000000..58a8c3d078f02a8
--- /dev/null
+++ b/llvm/test/Transforms/LoopVectorize/as_cast.ll
@@ -0,0 +1,40 @@
+; RUN: opt -passes=loop-vectorize %s -force-vector-width=1 -force-vector-interleave=2 -S -o - | FileCheck %s
+
+define void @foo(ptr addrspace(1) %in) {
+entry:
+  br label %loop
+
+loop:
+  %iter = phi i64 [ %next, %loop ], [ 0, %entry ]
+  %ascast = addrspacecast ptr addrspace(1) %in to ptr
+  %next = add i64 %iter, 1
+  %arrayidx = getelementptr inbounds i64, ptr %ascast, i64 %next
+  store i64 %next, ptr %arrayidx, align 4
+
+; check that we find the two interleaved blocks with ascast, gep and store:
+; CHECK: pred.store.if:
+; CHECK: [[ID1:%.*]] = add i64 %{{.*}}, 1
+; CHECK: [[AS1:%.*]] = addrspacecast ptr addrspace(1) %{{.*}} to ptr
+; CHECK: [[GEP1:%.*]] = getelementptr inbounds i64, ptr [[AS1]], i64 [[ID1]]
+; CHECK: store i64 [[ID1]], ptr [[GEP1]]
+
+; CHECK: pred.store.if1:
+; CHECK: [[ID2:%.*]] = add i64 %{{.*}}, 1
+; CHECK: [[AS2:%.*]] = addrspacecast ptr addrspace(1) %in to ptr
+; CHECK: [[GEP2:%.*]] = getelementptr inbounds i64, ptr [[AS2]], i64 [[ID2]]
+; CHECK: store i64 [[ID2]], ptr %9, align 4
+
+  %cmp = icmp eq i64 %next, 7
+  br i1 %cmp, label %exit, label %loop
+
+; check that we branch to the exit block
+; CHECK: middle.block:
+; CHECK: br i1 true, label %exit, label %scalar.ph
+
+exit:
+  ret void
+; CHECK: exit:
+; CHECK:   ret void
+}
+
+; CHECK: !{{[0-9]*}} = !{!"llvm.loop.isvectorized", i32 1}