[CIR] ArraySubscriptExpr on ExtVectorElementExpr (#169158)

Implement ArraySubscriptExpr support for ExtVectorElementExpr

Author: AmrDeveloper

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -683,6 +683,30 @@ RValue CIRGenFunction::emitLoadOfExtVectorElementLValue(LValue lv) {
   return RValue::get(resultVec);
 }
 
+/// Generates lvalue for partial ext_vector access.
+Address CIRGenFunction::emitExtVectorElementLValue(LValue lv,
+                                                   mlir::Location loc) {
+  Address vectorAddress = lv.getExtVectorAddress();
+  QualType elementTy = lv.getType()->castAs<VectorType>()->getElementType();
+  mlir::Type vectorElementTy = cgm.getTypes().convertType(elementTy);
+  Address castToPointerElement =
+      vectorAddress.withElementType(builder, vectorElementTy);
+
+  mlir::ArrayAttr extVecElts = lv.getExtVectorElts();
+  unsigned idx = getAccessedFieldNo(0, extVecElts);
+  mlir::Value idxValue =
+      builder.getConstInt(loc, mlir::cast<cir::IntType>(ptrDiffTy), idx);
+
+  mlir::Value elementValue = builder.getArrayElement(
+      loc, loc, castToPointerElement.getPointer(), vectorElementTy, idxValue,
+      /*shouldDecay=*/false);
+
+  const CharUnits eltSize = getContext().getTypeSizeInChars(elementTy);
+  const CharUnits alignment =
+      castToPointerElement.getAlignment().alignmentAtOffset(idx * eltSize);
+  return Address(elementValue, vectorElementTy, alignment);
+}
+
 static cir::FuncOp emitFunctionDeclPointer(CIRGenModule &cgm, GlobalDecl gd) {
   assert(!cir::MissingFeatures::weakRefReference());
   return cgm.getAddrOfFunction(gd);
@@ -1081,12 +1105,6 @@ static Address emitArraySubscriptPtr(CIRGenFunction &cgf,
 
 LValue
 CIRGenFunction::emitArraySubscriptExpr(const clang::ArraySubscriptExpr *e) {
-  if (isa<ExtVectorElementExpr>(e->getBase())) {
-    cgm.errorNYI(e->getSourceRange(),
-                 "emitArraySubscriptExpr: ExtVectorElementExpr");
-    return LValue::makeAddr(Address::invalid(), e->getType(), LValueBaseInfo());
-  }
-
   if (getContext().getAsVariableArrayType(e->getType())) {
     cgm.errorNYI(e->getSourceRange(),
                  "emitArraySubscriptExpr: VariableArrayType");
@@ -1116,15 +1134,30 @@ CIRGenFunction::emitArraySubscriptExpr(const clang::ArraySubscriptExpr *e) {
 
   // If the base is a vector type, then we are forming a vector element
   // with this subscript.
-  if (e->getBase()->getType()->isVectorType() &&
+  if (e->getBase()->getType()->isSubscriptableVectorType() &&
       !isa<ExtVectorElementExpr>(e->getBase())) {
     const mlir::Value idx = emitIdxAfterBase(/*promote=*/false);
-    const LValue lhs = emitLValue(e->getBase());
-    return LValue::makeVectorElt(lhs.getAddress(), idx, e->getBase()->getType(),
-                                 lhs.getBaseInfo());
+    const LValue lv = emitLValue(e->getBase());
+    return LValue::makeVectorElt(lv.getAddress(), idx, e->getBase()->getType(),
+                                 lv.getBaseInfo());
   }
 
   const mlir::Value idx = emitIdxAfterBase(/*promote=*/true);
+
+  // Handle the extvector case we ignored above.
+  if (isa<ExtVectorElementExpr>(e->getBase())) {
+    const LValue lv = emitLValue(e->getBase());
+    Address addr = emitExtVectorElementLValue(lv, cgm.getLoc(e->getExprLoc()));
+
+    QualType elementType = lv.getType()->castAs<VectorType>()->getElementType();
+    addr = emitArraySubscriptPtr(*this, cgm.getLoc(e->getBeginLoc()),
+                                 cgm.getLoc(e->getEndLoc()), addr, e->getType(),
+                                 idx, cgm.getLoc(e->getExprLoc()),
+                                 /*shouldDecay=*/false);
+
+    return makeAddrLValue(addr, elementType, lv.getBaseInfo());
+  }
+
   if (const Expr *array = getSimpleArrayDecayOperand(e->getBase())) {
     LValue arrayLV;
     if (const auto *ase = dyn_cast<ArraySubscriptExpr>(array))

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -531,6 +531,8 @@ class CIRGenFunction : public CIRGenTypeCache {
   /// that we can just remove the code.
   bool containsLabel(const clang::Stmt *s, bool ignoreCaseStmts = false);
 
+  Address emitExtVectorElementLValue(LValue lv, mlir::Location loc);
+
   class ConstantEmission {
     // Cannot use mlir::TypedAttr directly here because of bit availability.
     llvm::PointerIntPair<mlir::Attribute, 1, bool> valueAndIsReference;

clang/test/CIR/CodeGen/vector-ext-element.cpp

@@ -315,3 +315,27 @@ void element_expr_from_rvalue_with_vec_result() {
 // OGCG: %[[TMP_2:.*]] = load <4 x i32>, ptr %[[TMP_2_ADDR]], align 16
 // OGCG: %[[D_VALUE:.*]] = shufflevector <4 x i32> %[[TMP_2]], <4 x i32> poison, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
 // OGCG: store <4 x i32> %[[D_VALUE]], ptr %[[D_ADDR]], align 16
+
+void array_subscript_expr_with_element_expr_base() {
+  vi4 a;
+  a.xyz[1] = 2;
+}
+
+// CIR: %[[A_ADDR:.*]] = cir.alloca !cir.vector<4 x !s32i>, !cir.ptr<!cir.vector<4 x !s32i>>, ["a"]
+// CIR: %[[CONST_2:.*]] = cir.const #cir.int<2> : !s32i
+// CIR: %[[CONST_1:.*]] = cir.const #cir.int<1> : !s32i
+// CIR: %[[A_PTR:.*]] = cir.cast bitcast %0 : !cir.ptr<!cir.vector<4 x !s32i>> -> !cir.ptr<!s32i>
+// CIR: %[[CONST_0:.*]] = cir.const #cir.int<0> : !s64i
+// CIR: %[[VEC_MEMBER_EXPR:.*]] = cir.ptr_stride %[[A_PTR]], %[[CONST_0]] : (!cir.ptr<!s32i>, !s64i) -> !cir.ptr<!s32i>
+// CIR: %[[VEC_ELEM_PTR:.*]] = cir.ptr_stride %[[VEC_MEMBER_EXPR]], %[[CONST_1]] : (!cir.ptr<!s32i>, !s32i) -> !cir.ptr<!s32i>
+// CIR: cir.store {{.*}} %[[CONST_2]], %[[VEC_ELEM_PTR]] : !s32i, !cir.ptr<!s32i>
+
+// LLVM: %[[A_ADDR:.*]] = alloca <4 x i32>, i64 1, align 16
+// LLVM: %[[VEC_MEMBER_EXPR:.*]] = getelementptr i32, ptr %[[A_ADDR]], i64 0
+// LLVM: %[[VEC_ELEM_PTR:.*]] = getelementptr i32, ptr %[[VEC_MEMBER_EXPR]], i64 1
+// LLVM: store i32 2, ptr %[[VEC_ELEM_PTR]], align 4
+
+// OGCG: %[[A_ADDR:.*]] = alloca <4 x i32>, align 16
+// OGCG: %[[VEC_MEMBER_EXPR:.*]] = getelementptr inbounds i32, ptr %[[A_ADDR]], i64 0
+// OGCG: %[[VEC_ELEM_PTR:.*]] = getelementptr inbounds i32, ptr %[[VEC_MEMBER_EXPR]], i64 1
+// OGCG: store i32 2, ptr %[[VEC_ELEM_PTR]], align 4