[CIR] ExtVectorElementExpr with result Vector type

[AmrDeveloper]

Upstream ExtVectorElementExpr with result Vector type

[llvmbot]

@llvm/pr-subscribers-clang

@llvm/pr-subscribers-clangir

Author: Amr Hesham (AmrDeveloper)

Changes

Upstream ExtVectorElementExpr with result Vector type

---

Full diff: https://github.com/llvm/llvm-project/pull/167925.diff

3 Files Affected:

• (modified) clang/lib/CIR/CodeGen/CIRGenBuilder.h (+28)
• (modified) clang/lib/CIR/CodeGen/CIRGenExpr.cpp (+13-3)
• (modified) clang/test/CIR/CodeGen/vector-ext-element.cpp (+39)

diff --git a/clang/lib/CIR/CodeGen/CIRGenBuilder.h b/clang/lib/CIR/CodeGen/CIRGenBuilder.h
index e5066fac19185..d7177b04623ef 100644
--- a/clang/lib/CIR/CodeGen/CIRGenBuilder.h
+++ b/clang/lib/CIR/CodeGen/CIRGenBuilder.h
@@ -572,6 +572,34 @@ class CIRGenBuilderTy : public cir::CIRBaseBuilderTy {
                                       info.isSigned, isLvalueVolatile,
                                       addr.getAlignment().getAsAlign().value());
   }
+
+  cir::VecShuffleOp
+  createVecShuffle(mlir::Location loc, mlir::Value vec1, mlir::Value vec2,
+                   llvm::ArrayRef<mlir::Attribute> maskAttrs) {
+    auto vecType = mlir::cast<cir::VectorType>(vec1.getType());
+    auto resultTy = cir::VectorType::get(getContext(), vecType.getElementType(),
+                                         maskAttrs.size());
+    return cir::VecShuffleOp::create(*this, loc, resultTy, vec1, vec2,
+                                     getArrayAttr(maskAttrs));
+  }
+
+  cir::VecShuffleOp createVecShuffle(mlir::Location loc, mlir::Value vec1,
+                                     mlir::Value vec2,
+                                     llvm::ArrayRef<int64_t> mask) {
+    llvm::SmallVector<mlir::Attribute> maskAttrs;
+    for (int32_t idx : mask)
+      maskAttrs.push_back(cir::IntAttr::get(getSInt32Ty(), idx));
+    return createVecShuffle(loc, vec1, vec2, maskAttrs);
+  }
+
+  cir::VecShuffleOp createVecShuffle(mlir::Location loc, mlir::Value vec1,
+                                     llvm::ArrayRef<int64_t> mask) {
+    /// Create a unary shuffle. The second vector operand of the IR instruction
+    /// is poison.
+    cir::ConstantOp poison =
+        getConstant(loc, cir::PoisonAttr::get(vec1.getType()));
+    return createVecShuffle(loc, vec1, poison, mask);
+  }
 };
 
 } // namespace clang::CIRGen
diff --git a/clang/lib/CIR/CodeGen/CIRGenExpr.cpp b/clang/lib/CIR/CodeGen/CIRGenExpr.cpp
index c67493a913d58..9c6d26546ac7e 100644
--- a/clang/lib/CIR/CodeGen/CIRGenExpr.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenExpr.cpp
@@ -669,9 +669,19 @@ RValue CIRGenFunction::emitLoadOfExtVectorElementLValue(LValue lv) {
     return RValue::get(cir::VecExtractOp::create(builder, loc, vec, index));
   }
 
-  cgm.errorNYI(
-      loc, "emitLoadOfExtVectorElementLValue: Result of expr is vector type");
-  return {};
+  // Always use shuffle vector to try to retain the original program structure
+  const unsigned numResultElts = exprVecTy->getNumElements();
+  SmallVector<int64_t> mask;
+  for (unsigned i = 0; i != numResultElts; ++i)
+    mask.push_back(getAccessedFieldNo(i, elts));
+
+  cir::VecShuffleOp resultVec = builder.createVecShuffle(loc, vec, mask);
+  if (lv.getType()->isExtVectorBoolType()) {
+    cgm.errorNYI(loc, "emitLoadOfExtVectorElementLValue: ExtVectorBoolType");
+    return {};
+  }
+
+  return RValue::get(resultVec);
 }
 
 static cir::FuncOp emitFunctionDeclPointer(CIRGenModule &cgm, GlobalDecl gd) {
diff --git a/clang/test/CIR/CodeGen/vector-ext-element.cpp b/clang/test/CIR/CodeGen/vector-ext-element.cpp
index de9d53936d2eb..e724eb58ff9aa 100644
--- a/clang/test/CIR/CodeGen/vector-ext-element.cpp
+++ b/clang/test/CIR/CodeGen/vector-ext-element.cpp
@@ -5,6 +5,7 @@
 // RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -Wno-unused-value -emit-llvm %s -o %t.ll
 // RUN: FileCheck --input-file=%t.ll %s -check-prefix=OGCG
 
+typedef int vi2 __attribute__((ext_vector_type(2)));
 typedef int vi4 __attribute__((ext_vector_type(4)));
 
 void element_expr_from_gl() {
@@ -44,3 +45,41 @@ void element_expr_from_gl() {
 // OGCG: %[[TMP_A:.*]] = load <4 x i32>, ptr %[[A_ADDR]], align 16
 // OGCG: %[[ELEM_1:.*]] = extractelement <4 x i32> %[[TMP_A]], i64 1
 // OGCG: store i32 %[[ELEM_1]], ptr %[[Y_ADDR]], align 4
+
+void element_expr_from_gl_with_vec_result() {
+  vi4 a;
+  vi2 b = a.xy;
+  vi4 c = a.wzyx;
+}
+
+// CIR: %[[A_ADDR:.*]] = cir.alloca !cir.vector<4 x !s32i>, !cir.ptr<!cir.vector<4 x !s32i>>, ["a"]
+// CIR: %[[B_ADDR:.*]] = cir.alloca !cir.vector<2 x !s32i>, !cir.ptr<!cir.vector<2 x !s32i>>, ["b", init]
+// CIR: %[[C_ADDR:.*]] = cir.alloca !cir.vector<4 x !s32i>, !cir.ptr<!cir.vector<4 x !s32i>>, ["c", init]
+// CIR: %[[TMP_A:.*]] = cir.load {{.*}} %[[A_ADDR]] : !cir.ptr<!cir.vector<4 x !s32i>>, !cir.vector<4 x !s32i>
+// CIR: %[[POISON:.*]] = cir.const #cir.poison : !cir.vector<4 x !s32i>
+// CIR: %[[B_VALUE:.*]] = cir.vec.shuffle(%[[TMP_A]], %[[POISON]] : !cir.vector<4 x !s32i>) [#cir.int<0> : !s32i, #cir.int<1> : !s32i] : !cir.vector<2 x !s32i>
+// CIR: cir.store {{.*}} %[[B_VALUE]], %[[B_ADDR]] : !cir.vector<2 x !s32i>, !cir.ptr<!cir.vector<2 x !s32i>>
+// CIR: %[[TMP_A:.*]] = cir.load {{.*}} %[[A_ADDR]] : !cir.ptr<!cir.vector<4 x !s32i>>, !cir.vector<4 x !s32i>
+// CIR: %[[POISON:.*]] = cir.const #cir.poison : !cir.vector<4 x !s32i>
+// CIR: %[[C_VALUE:.*]] = cir.vec.shuffle(%[[TMP_A]], %[[POISON]] : !cir.vector<4 x !s32i>) [#cir.int<3> : !s32i, #cir.int<2> : !s32i, #cir.int<1> : !s32i, #cir.int<0> : !s32i] : !cir.vector<4 x !s32i>
+// CIR: cir.store {{.*}} %[[C_VALUE]], %[[C_ADDR]] : !cir.vector<4 x !s32i>, !cir.ptr<!cir.vector<4 x !s32i>>
+
+// LLVM: %[[A_ADDR:.*]] = alloca <4 x i32>, i64 1, align 16
+// LLVM: %[[B_ADDR:.*]] = alloca <2 x i32>, i64 1, align 8
+// LLVM: %[[C_ADDR:.*]] = alloca <4 x i32>, i64 1, align 16
+// LLVM: %[[TMP_A:.*]] = load <4 x i32>, ptr %[[A_ADDR]], align 16
+// LLVM: %[[B_VALUE:.*]] = shufflevector <4 x i32> %[[TMP_A]], <4 x i32> poison, <2 x i32> <i32 0, i32 1>
+// LLVM: store <2 x i32> %[[B_VALUE]], ptr %[[B_ADDR]], align 8
+// LLVM: %[[TMP_A:.*]] = load <4 x i32>, ptr %[[A_ADDR]], align 16
+// LLVM: %[[C_VALUE:.*]] = shufflevector <4 x i32> %[[TMP_A]], <4 x i32> poison, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+// LLVM: store <4 x i32> %[[C_VALUE]], ptr %[[C_ADDR]], align 16
+
+// OGCG: %[[A_ADDR:.*]] = alloca <4 x i32>, align 16
+// OGCG: %[[B_ADDR:.*]] = alloca <2 x i32>, align 8
+// OGCG: %[[C_ADDR:.*]] = alloca <4 x i32>, align 16
+// OGCG: %[[TMP_A:.*]] = load <4 x i32>, ptr %[[A_ADDR]], align 16
+// OGCG: %[[B_VALUE:.*]] = shufflevector <4 x i32> %[[TMP_A]], <4 x i32> poison, <2 x i32> <i32 0, i32 1>
+// OGCG: store <2 x i32> %[[B_VALUE]], ptr %[[B_ADDR]], align 8
+// OGCG: %[[TMP_A:.*]] = load <4 x i32>, ptr %[[A_ADDR]], align 16
+// OGCG: %[[C_VALUE:.*]] = shufflevector <4 x i32> %[[TMP_A]], <4 x i32> poison, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
+// OGCG: store <4 x i32> %[[C_VALUE]], ptr %[[C_ADDR]], align 16

[xlauko]

Suggested change

	llvm::SmallVector<mlir::Attribute> maskAttrs;
	for (int32_t idx : mask)
	maskAttrs.push_back(cir::IntAttr::get(getSInt32Ty(), idx));
	auto maskAttrs =
	llvm::to_vector<llvm::SmallVector<mlir::Attribute>>(
	llvm::map_range(mask, [&](int32_t idx) {
	return cir::getSInt32(idx);
	}));

[AmrDeveloper]

In the source code, I found that to_vector takes size, not type. Did you mean to_vector_of 🤔 ?

  auto maskAttrs = llvm::to_vector_of<mlir::Attribute>(
        llvm::map_range(mask, [&](int32_t idx) {
          return cir::IntAttr::get(getSInt32Ty(), idx);
        }));

[xlauko]

lgtm % nits