diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td index 237daed32532a..7db6725afb80c 100644 --- a/clang/include/clang/CIR/Dialect/IR/CIROps.td +++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td @@ -2192,6 +2192,7 @@ def VecShuffleDynamicOp : CIR_Op<"vec.shuffle.dynamic", }]; let hasVerifier = 1; + let hasFolder = 1; } #endif // CLANG_CIR_DIALECT_IR_CIROPS_TD diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp index 36f050de9f8bb..26ebe6c1dbcd1 100644 --- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp +++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp @@ -21,6 +21,7 @@ #include "clang/CIR/Dialect/IR/CIROpsDialect.cpp.inc" #include "clang/CIR/Dialect/IR/CIROpsEnums.cpp.inc" #include "clang/CIR/MissingFeatures.h" + #include using namespace mlir; @@ -1579,6 +1580,37 @@ OpFoldResult cir::VecExtractOp::fold(FoldAdaptor adaptor) { // VecShuffleDynamicOp //===----------------------------------------------------------------------===// +OpFoldResult cir::VecShuffleDynamicOp::fold(FoldAdaptor adaptor) { + mlir::Attribute vec = adaptor.getVec(); + mlir::Attribute indices = adaptor.getIndices(); + if (mlir::isa_and_nonnull(vec) && + mlir::isa_and_nonnull(indices)) { + auto vecAttr = mlir::cast(vec); + auto indicesAttr = mlir::cast(indices); + auto vecTy = mlir::cast(vecAttr.getType()); + + mlir::ArrayAttr vecElts = vecAttr.getElts(); + mlir::ArrayAttr indicesElts = indicesAttr.getElts(); + + const uint64_t numElements = vecElts.size(); + + SmallVector elements; + elements.reserve(numElements); + + const uint64_t maskBits = llvm::NextPowerOf2(numElements - 1) - 1; + for (const auto &idxAttr : indicesElts.getAsRange()) { + uint64_t idxValue = idxAttr.getUInt(); + uint64_t newIdx = idxValue & maskBits; + elements.push_back(vecElts[newIdx]); + } + + return cir::ConstVectorAttr::get( + vecTy, mlir::ArrayAttr::get(getContext(), elements)); + } + + return {}; +} + LogicalResult cir::VecShuffleDynamicOp::verify() { // The number of elements in the two input vectors must match. if (getVec().getType().getSize() != diff --git a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp index fb000adee04c6..7d03e374c27e8 100644 --- a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp +++ b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp @@ -138,10 +138,10 @@ void CIRCanonicalizePass::runOnOperation() { assert(!cir::MissingFeatures::complexRealOp()); assert(!cir::MissingFeatures::complexImagOp()); assert(!cir::MissingFeatures::callOp()); - // CastOp, UnaryOp and VecExtractOp are here to perform a manual `fold` in - // applyOpPatternsGreedily. + // CastOp, UnaryOp, VecExtractOp and VecShuffleDynamicOp are here to perform + // a manual `fold` in applyOpPatternsGreedily. if (isa(op)) + VecExtractOp, VecShuffleDynamicOp>(op)) ops.push_back(op); }); diff --git a/clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir b/clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir new file mode 100644 index 0000000000000..46ab04502afec --- /dev/null +++ b/clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir @@ -0,0 +1,31 @@ +// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s + +!s32i = !cir.int + +module { + cir.func @fold_shuffle_dynamic_vector_op_test() -> !cir.vector<4 x !s32i> { + %vec = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i, #cir.int<4> : !s32i]> : !cir.vector<4 x !s32i> + %indices = cir.const #cir.const_vector<[#cir.int<8> : !s32i, #cir.int<7> : !s32i, #cir.int<6> : !s32i, #cir.int<5> : !s32i]> : !cir.vector<4 x !s32i> + %new_vec = cir.vec.shuffle.dynamic %vec : !cir.vector<4 x !s32i>, %indices : !cir.vector<4 x !s32i> + cir.return %new_vec : !cir.vector<4 x !s32i> + } + + // Masking indices [8, 7, 6, 5] AND 3 = [0, 3, 2, 1] + // CHECK: cir.func @fold_shuffle_dynamic_vector_op_test() -> !cir.vector<4 x !s32i> { + // CHECK-NEXT: %[[NEW_VEC:.*]] = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<4> : !s32i, #cir.int<3> : !s32i, #cir.int<2> : !s32i]> : !cir.vector<4 x !s32i> + // CHECK-NEXT: cir.return %[[NEW_VEC:.*]] : !cir.vector<4 x !s32i> + + cir.func @fold_shuffle_dynamic_vector_op_test_2() -> !cir.vector<4 x !s32i> { + %vec = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i, #cir.int<4> : !s32i]> : !cir.vector<4 x !s32i> + %indices = cir.const #cir.const_vector<[#cir.int<3> : !s32i, #cir.int<2> : !s32i, #cir.int<1> : !s32i, #cir.int<0> : !s32i]> : !cir.vector<4 x !s32i> + %new_vec = cir.vec.shuffle.dynamic %vec : !cir.vector<4 x !s32i>, %indices : !cir.vector<4 x !s32i> + cir.return %new_vec : !cir.vector<4 x !s32i> + } + + // Masking indices [3, 2, 1, 0] AND 3 = [3, 2, 1, 0] + // CHECK: cir.func @fold_shuffle_dynamic_vector_op_test_2() -> !cir.vector<4 x !s32i> { + // CHECK-NEXT: %[[NEW_VEC:.*]] = cir.const #cir.const_vector<[#cir.int<4> : !s32i, #cir.int<3> : !s32i, #cir.int<2> : !s32i, #cir.int<1> : !s32i]> : !cir.vector<4 x !s32i> + // CHECK-NEXT: cir.return %[[NEW_VEC:.*]] : !cir.vector<4 x !s32i> +} + +