From 6f215ca680fdda6eff0114bbd9d26a3efd8b21be Mon Sep 17 00:00:00 2001
From: Paul Walker <paul.walker@arm.com>
Date: Sun, 22 May 2022 13:49:01 +0100
Subject: [PATCH] [SelectionDAG] Add support to widen ISD::STEP_VECTOR
 operations.

Fixes: #55165

Differential Revision: https://reviews.llvm.org/D126168
---
 .../CodeGen/SelectionDAG/LegalizeVectorTypes.cpp   |  1 +
 llvm/test/CodeGen/AArch64/sve-stepvector.ll        | 26 ++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index ffb36ee..69974df 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -3615,6 +3615,7 @@ void DAGTypeLegalizer::WidenVectorResult(SDNode *N, unsigned ResNo) {
   case ISD::EXTRACT_SUBVECTOR: Res = WidenVecRes_EXTRACT_SUBVECTOR(N); break;
   case ISD::INSERT_VECTOR_ELT: Res = WidenVecRes_INSERT_VECTOR_ELT(N); break;
   case ISD::LOAD:              Res = WidenVecRes_LOAD(N); break;
+  case ISD::STEP_VECTOR:
   case ISD::SPLAT_VECTOR:
   case ISD::SCALAR_TO_VECTOR:
     Res = WidenVecRes_ScalarOp(N);
diff --git a/llvm/test/CodeGen/AArch64/sve-stepvector.ll b/llvm/test/CodeGen/AArch64/sve-stepvector.ll
index 61aafb6..4b3cb9c 100644
--- a/llvm/test/CodeGen/AArch64/sve-stepvector.ll
+++ b/llvm/test/CodeGen/AArch64/sve-stepvector.ll
@@ -45,6 +45,20 @@ entry:
 
 ; ILLEGAL INTEGER TYPES
 
+define <vscale x 6 x i64> @stepvector_nxv6i64() {
+; CHECK-LABEL: stepvector_nxv6i64:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z0.d, #0, #1
+; CHECK-NEXT:    mov z1.d, z0.d
+; CHECK-NEXT:    mov z2.d, z0.d
+; CHECK-NEXT:    incd z1.d
+; CHECK-NEXT:    incd z2.d, all, mul #2
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 6 x i64> @llvm.experimental.stepvector.nxv6i64()
+  ret <vscale x 6 x i64> %0
+}
+
 define <vscale x 4 x i64> @stepvector_nxv4i64() {
 ; CHECK-LABEL: stepvector_nxv4i64:
 ; CHECK:       // %bb.0: // %entry
@@ -73,6 +87,16 @@ entry:
   ret <vscale x 16 x i32> %0
 }
 
+define <vscale x 3 x i32> @stepvector_nxv3i32() {
+; CHECK-LABEL: stepvector_nxv3i32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z0.s, #0, #1
+; CHECK-NEXT:    ret
+entry:
+  %0 = call <vscale x 3 x i32> @llvm.experimental.stepvector.nxv3i32()
+  ret <vscale x 3 x i32> %0
+}
+
 define <vscale x 2 x i32> @stepvector_nxv2i32() {
 ; CHECK-LABEL: stepvector_nxv2i32:
 ; CHECK:       // %bb.0: // %entry
@@ -422,8 +446,10 @@ declare <vscale x 4 x i32> @llvm.experimental.stepvector.nxv4i32()
 declare <vscale x 8 x i16> @llvm.experimental.stepvector.nxv8i16()
 declare <vscale x 16 x i8> @llvm.experimental.stepvector.nxv16i8()
 
+declare <vscale x 6 x i64> @llvm.experimental.stepvector.nxv6i64()
 declare <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
 declare <vscale x 16 x i32> @llvm.experimental.stepvector.nxv16i32()
+declare <vscale x 3 x i32> @llvm.experimental.stepvector.nxv3i32()
 declare <vscale x 2 x i32> @llvm.experimental.stepvector.nxv2i32()
 declare <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
 declare <vscale x 4 x i16> @llvm.experimental.stepvector.nxv4i16()
-- 
2.7.4