[clang] [Clang][AArch64][SVE] Allow write to SVE vector elements using the subscript operator (PR #91965)

[Momchil Velikov via cfe-commits]

https://github.com/momchil-velikov updated https://github.com/llvm/llvm-project/pull/91965

>From b1b69ffcaf4525a66dde1ae7f1a022c85204a579 Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov at arm.com>
Date: Mon, 20 May 2024 16:25:43 +0100
Subject: [PATCH 1/2] [Clang][AArch64][SVE] Allow write to SVE vector elements
 using the subscript operator

The patch at https://reviews.llvm.org/D122732 introduced using
the array subscript operator for SVE vectors, however it also
causes an ICE when the subscripting expression is used
as an lvalue.

This patches fixes the error. Lvalue subscripting expressions are
emitted as LLVM IR `insertvector`.

Change-Id: I46d0333d8ed8508cd9cd23e02dd1c2d48fb74cd2
---
 clang/lib/CodeGen/CGExpr.cpp                          |  2 +-
 clang/lib/Sema/SemaExpr.cpp                           |  2 +-
 clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c | 10 ++++++++++
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index cd1c48b420382..6f9237e2067f5 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -4180,7 +4180,7 @@ LValue CodeGenFunction::EmitArraySubscriptExpr(const ArraySubscriptExpr *E,
 
   // If the base is a vector type, then we are forming a vector element lvalue
   // with this subscript.
-  if (E->getBase()->getType()->isVectorType() &&
+  if (E->getBase()->getType()->isSubscriptableVectorType() &&
       !isa<ExtVectorElementExpr>(E->getBase())) {
     // Emit the vector as an lvalue to get its address.
     LValue LHS = EmitLValue(E->getBase());
diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
index 326879b0883fa..49541edf106e1 100644
--- a/clang/lib/Sema/SemaExpr.cpp
+++ b/clang/lib/Sema/SemaExpr.cpp
@@ -5185,7 +5185,7 @@ Sema::CreateBuiltinArraySubscriptExpr(Expr *Base, SourceLocation LLoc,
   }
 
   // Perform default conversions.
-  if (!LHSExp->getType()->getAs<VectorType>()) {
+  if (!LHSExp->getType()->isSubscriptableVectorType()) {
     ExprResult Result = DefaultFunctionArrayLvalueConversion(LHSExp);
     if (Result.isInvalid())
       return ExprError();
diff --git a/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c b/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c
index fb60c6d100ce6..634423765c4c3 100644
--- a/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c
+++ b/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c
@@ -88,3 +88,13 @@ float subscript_float32(svfloat32_t a, size_t b) {
 double subscript_float64(svfloat64_t a, size_t b) {
   return a[b];
 }
+
+// CHECK-LABEL: @subscript_write_float32(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[VECINS:%.*]] = insertelement <vscale x 4 x float> [[A:%.*]], float 1.000000e+00, i64 [[B:%.*]]
+// CHECK-NEXT:    ret <vscale x 4 x float> [[VECINS]]
+//
+svfloat32_t subscript_write_float32(svfloat32_t a, size_t b) {
+  a[b] = 1.0f;
+  return a;
+}

>From 08a155b49a6b9c859ba8569170e0f71e63b76735 Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov at arm.com>
Date: Mon, 20 May 2024 16:26:06 +0100
Subject: [PATCH 2/2] [fixiup] Add a test using compound assignment operator

Change-Id: I81e1fd4f23eb65a96e71015de7a4562fcbc53c0f
---
 .../test/CodeGen/aarch64-sve-vector-subscript-ops.c  | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c b/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c
index 634423765c4c3..52a05d010de9b 100644
--- a/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c
+++ b/clang/test/CodeGen/aarch64-sve-vector-subscript-ops.c
@@ -98,3 +98,15 @@ svfloat32_t subscript_write_float32(svfloat32_t a, size_t b) {
   a[b] = 1.0f;
   return a;
 }
+
+// CHECK-LABEL: @subscript_read_write_float32(
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[VECEXT:%.*]] = extractelement <vscale x 4 x float> [[A:%.*]], i64 [[B:%.*]]
+// CHECK-NEXT:    [[ADD:%.*]] = fadd float [[VECEXT]], 1.000000e+00
+// CHECK-NEXT:    [[VECINS:%.*]] = insertelement <vscale x 4 x float> [[A]], float [[ADD]], i64 [[B]]
+// CHECK-NEXT:    ret <vscale x 4 x float> [[VECINS]]
+//
+svfloat32_t subscript_read_write_float32(svfloat32_t a, size_t b) {
+  a[b] += 1.0f;
+  return a;
+}