Skip to content

[IR2Vec] Add out-of-place arithmetic operators to Embedding class #145118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions llvm/include/llvm/Analysis/IR2Vec.h
Original file line number Diff line number Diff line change
Expand Up @@ -107,9 +107,12 @@ struct Embedding {
const std::vector<double> &getData() const { return Data; }

/// Arithmetic operators
Embedding &operator+=(const Embedding &RHS);
Embedding &operator-=(const Embedding &RHS);
Embedding &operator*=(double Factor);
LLVM_ABI Embedding &operator+=(const Embedding &RHS);
LLVM_ABI Embedding operator+(const Embedding &RHS) const;
LLVM_ABI Embedding &operator-=(const Embedding &RHS);
LLVM_ABI Embedding operator-(const Embedding &RHS) const;
LLVM_ABI Embedding &operator*=(double Factor);
LLVM_ABI Embedding operator*(double Factor) const;

/// Adds Src Embedding scaled by Factor with the called Embedding.
/// Called_Embedding += Src * Factor
Expand Down
19 changes: 18 additions & 1 deletion llvm/lib/Analysis/IR2Vec.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -70,27 +70,44 @@ inline bool fromJSON(const llvm::json::Value &E, Embedding &Out,
// ==----------------------------------------------------------------------===//
// Embedding
//===----------------------------------------------------------------------===//

Embedding &Embedding::operator+=(const Embedding &RHS) {
assert(this->size() == RHS.size() && "Vectors must have the same dimension");
std::transform(this->begin(), this->end(), RHS.begin(), this->begin(),
std::plus<double>());
return *this;
}

Embedding Embedding::operator+(const Embedding &RHS) const {
Embedding Result(*this);
Result += RHS;
return Result;
}

Embedding &Embedding::operator-=(const Embedding &RHS) {
assert(this->size() == RHS.size() && "Vectors must have the same dimension");
std::transform(this->begin(), this->end(), RHS.begin(), this->begin(),
std::minus<double>());
return *this;
}

Embedding Embedding::operator-(const Embedding &RHS) const {
Embedding Result(*this);
Result -= RHS;
return Result;
}

Embedding &Embedding::operator*=(double Factor) {
std::transform(this->begin(), this->end(), this->begin(),
[Factor](double Elem) { return Elem * Factor; });
return *this;
}

Embedding Embedding::operator*(double Factor) const {
Embedding Result(*this);
Result *= Factor;
return Result;
}

Embedding &Embedding::scaleAndAdd(const Embedding &Src, float Factor) {
assert(this->size() == Src.size() && "Vectors must have the same dimension");
for (size_t Itr = 0; Itr < this->size(); ++Itr)
Expand Down
39 changes: 39 additions & 0 deletions llvm/unittests/Analysis/IR2VecTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,18 @@ TEST(EmbeddingTest, ConstructorsAndAccessors) {
}
}

TEST(EmbeddingTest, AddVectorsOutOfPlace) {
Embedding E1 = {1.0, 2.0, 3.0};
Embedding E2 = {0.5, 1.5, -1.0};

Embedding E3 = E1 + E2;
EXPECT_THAT(E3, ElementsAre(1.5, 3.5, 2.0));

// Check that E1 and E2 are unchanged
EXPECT_THAT(E1, ElementsAre(1.0, 2.0, 3.0));
EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));
}

TEST(EmbeddingTest, AddVectors) {
Embedding E1 = {1.0, 2.0, 3.0};
Embedding E2 = {0.5, 1.5, -1.0};
Expand All @@ -120,6 +132,18 @@ TEST(EmbeddingTest, AddVectors) {
EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));
}

TEST(EmbeddingTest, SubtractVectorsOutOfPlace) {
Embedding E1 = {1.0, 2.0, 3.0};
Embedding E2 = {0.5, 1.5, -1.0};

Embedding E3 = E1 - E2;
EXPECT_THAT(E3, ElementsAre(0.5, 0.5, 4.0));

// Check that E1 and E2 are unchanged
EXPECT_THAT(E1, ElementsAre(1.0, 2.0, 3.0));
EXPECT_THAT(E2, ElementsAre(0.5, 1.5, -1.0));
}

TEST(EmbeddingTest, SubtractVectors) {
Embedding E1 = {1.0, 2.0, 3.0};
Embedding E2 = {0.5, 1.5, -1.0};
Expand All @@ -137,6 +161,15 @@ TEST(EmbeddingTest, ScaleVector) {
EXPECT_THAT(E1, ElementsAre(0.5, 1.0, 1.5));
}

TEST(EmbeddingTest, ScaleVectorOutOfPlace) {
Embedding E1 = {1.0, 2.0, 3.0};
Embedding E2 = E1 * 0.5f;
EXPECT_THAT(E2, ElementsAre(0.5, 1.0, 1.5));

// Check that E1 is unchanged
EXPECT_THAT(E1, ElementsAre(1.0, 2.0, 3.0));
}

TEST(EmbeddingTest, AddScaledVector) {
Embedding E1 = {1.0, 2.0, 3.0};
Embedding E2 = {2.0, 0.5, -1.0};
Expand Down Expand Up @@ -180,6 +213,12 @@ TEST(EmbeddingTest, AccessOutOfBounds) {
EXPECT_DEATH(E[4] = 4.0, "Index out of bounds");
}

TEST(EmbeddingTest, MismatchedDimensionsAddVectorsOutOfPlace) {
Embedding E1 = {1.0, 2.0};
Embedding E2 = {1.0};
EXPECT_DEATH(E1 + E2, "Vectors must have the same dimension");
}

TEST(EmbeddingTest, MismatchedDimensionsAddVectors) {
Embedding E1 = {1.0, 2.0};
Embedding E2 = {1.0};
Expand Down
Loading