From b6f34bbe67fe9006c20c55bba33f55a18a212114 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Mon, 18 Sep 2023 14:32:52 -0400 Subject: [PATCH 1/5] Initial commit --- csrc/evaluator_common.cpp | 99 +++++++++++++++++++++++++++++++++++++++ csrc/evaluator_common.h | 27 +++++++++-- 2 files changed, 121 insertions(+), 5 deletions(-) diff --git a/csrc/evaluator_common.cpp b/csrc/evaluator_common.cpp index e19631a59aa..58e9314512c 100644 --- a/csrc/evaluator_common.cpp +++ b/csrc/evaluator_common.cpp @@ -353,6 +353,14 @@ NaiveValueMachine::NaiveValueMachine(PrecomputedValues& precomputed_values) makeUnaryOp(uop); } else if (auto bop = dynamic_cast(def)) { makeBinaryOp(bop); + } else if (auto top = dynamic_cast(def)) { + makeTernaryOp(top); + } else if (auto lsop = dynamic_cast(def)) { + NVF_ERROR( + lsop->opType() == LoadStoreOpType::Set, + "NaiveValueMachine: unsupported LoadStoreOpType: ", + lsop->opType()); + makeSetOp(lsop); } else { // There could be some ops not supported yet. For these ops, we will // bind their outputs. So ignoring them here. @@ -435,14 +443,49 @@ void NaiveValueMachine::makeBinaryOp(BinaryOp* bop) { dest_[index] = out; } +void NaiveValueMachine::makeTernaryOp(TernaryOp* top) { + int in0 = top->inputs()[0]->evaluatorIndex(); + int in1 = top->inputs()[1]->evaluatorIndex(); + int in2 = top->inputs()[2]->evaluatorIndex(); + int out = top->outputs()[0]->evaluatorIndex(); + + NVF_ERROR(in0 >= 0, "Integer Machine: unknown first input: ", top); + NVF_ERROR(in1 >= 0, "Integer Machine: unknown second input: ", top); + NVF_ERROR(in2 >= 0, "Integer Machine: unknown third input: ", top); + NVF_ERROR(out >= 0, "Integer Machine: unknown out: ", top); + + int index = makeInstructionEntry(); + inst_type_[index] = InstructionType::TERNARY_OP; + top_type_[index] = top->getTernaryOpType(); + src0_[index] = in0; + src1_[index] = in1; + src2_[index] = in2; + dest_[index] = out; +} + +void NaiveValueMachine::makeSetOp(LoadStoreOp* lsop) { + int in = lsop->in()->evaluatorIndex(); + int out = lsop->out()->evaluatorIndex(); + + NVF_ERROR(in >= 0, "Integer Machine: unknown input: ", lsop); + NVF_ERROR(out >= 0, "Integer Machine: unknown out: ", lsop); + + int index = makeInstructionEntry(); + inst_type_[index] = InstructionType::SET_OP; + src0_[index] = in; + dest_[index] = out; +} + int NaiveValueMachine::makeInstructionEntry() { int index = num_of_instructions_++; inst_type_.emplace_back(InstructionType::UNARY_OP); uop_type_.emplace_back(UnaryOpType::Abs); bop_type_.emplace_back(BinaryOpType::Add); + top_type_.emplace_back(TernaryOpType::Where); data_type_.emplace_back(DataType::Null); src0_.emplace_back(-1); src1_.emplace_back(-1); + src2_.emplace_back(-1); dest_.emplace_back(-1); return index; } @@ -459,6 +502,9 @@ void NaiveValueMachine::runInstruction(int index) { case InstructionType::BINARY_OP: runBinaryOp(index); break; + case InstructionType::TERNARY_OP: + runTernaryOp(index); + break; } } @@ -574,6 +620,59 @@ void NaiveValueMachine::runBinaryOp(int index) { case BinaryOpType::Gcd: dest = gcd(lhs, rhs); break; + case BinaryOpType::LT: + dest = lhs < rhs; + break; + case BinaryOpType::LE: + dest = lhs <= rhs; + break; + case BinaryOpType::Eq: + dest = lhs == rhs; + break; + case BinaryOpType::NE: + dest = lhs != rhs; + break; + case BinaryOpType::GE: + dest = lhs >= rhs; + break; + case BinaryOpType::GT: + dest = lhs > rhs; + break; + default: + NVF_CHECK(false, "Unexpected operator type ", bop_type_[index]); + } + + precomputed_values_.defined_[dest_index] = true; +} + +void NaiveValueMachine::runTernaryOp(int index) { + using namespace PolymorphicValue_functions; + int src0_index = src0_[index]; + int src1_index = src1_[index]; + int src2_index = src2_[index]; + bool src0_is_const = precomputed_values_.is_constant_[src0_index]; + bool src1_is_const = precomputed_values_.is_constant_[src1_index]; + bool src2_is_const = precomputed_values_.is_constant_[src2_index]; + + bool src_defined = + (precomputed_values_.defined_[src0_index] || src0_is_const) && + (precomputed_values_.defined_[src1_index] || src1_is_const) && + (precomputed_values_.defined_[src2_index] || src2_is_const); + + if (!src_defined) { + return; + } + int dest_index = dest_[index]; + + auto& a = precomputed_values_.values_[src0_index]; + auto& b = precomputed_values_.values_[src1_index]; + auto& c = precomputed_values_.values_[src2_index]; + auto& dest = precomputed_values_.values_[dest_index]; + + switch (top_type_[index]) { + case TernaryOpType::Where: + dest = a ? b : c; + break; default: NVF_CHECK(!"Unexpected operator type"); } diff --git a/csrc/evaluator_common.h b/csrc/evaluator_common.h index c8af9d34865..28b4a1e1d9f 100644 --- a/csrc/evaluator_common.h +++ b/csrc/evaluator_common.h @@ -30,9 +30,8 @@ struct TensorArgAbstract; //! PrecomputedValues that will provide the workspace //! containing the concrete values for the values. class NaiveValueMachine { - //! The generic types of instructions supported for this - //! machine, currently only binary and unary. - enum class InstructionType { UNARY_OP, BINARY_OP, SET_OP }; + //! The generic types of instructions supported for this machine. + enum class InstructionType { UNARY_OP, BINARY_OP, TERNARY_OP, SET_OP }; public: //! Constructor lowers all the expr IR nodes stored in precomputed_values @@ -56,6 +55,13 @@ class NaiveValueMachine { //! Convert an binary IR expr to an instruction void makeBinaryOp(BinaryOp* bop); + //! Convert an ternary IR expr to an instruction + void makeTernaryOp(TernaryOp* bop); + + //! Convert a LoadStoreOp expr to an instruction. This assumes lsop->opType() + //! is equal to LoadStoreOpType::Set. + void makeSetOp(LoadStoreOp* lsop); + //! Create an empty instruction with all default values //! and place it at the end of the instruction buffer. int makeInstructionEntry(); @@ -71,6 +77,9 @@ class NaiveValueMachine { //! Runs a binary operation at given index of instruction buffer void runBinaryOp(int index); + //! Runs a ternary operation at given index of instruction buffer + void runTernaryOp(int index); + private: friend PrecomputedValues; @@ -98,10 +107,14 @@ class NaiveValueMachine { //! value at each index corresponding other ops. std::vector data_type_; - //! Unary operator type if applicable, contains a default - //! value at each index corresponding to a unary op. + //! Binary operator type if applicable, contains a default + //! value at each index corresponding to a binary op. std::vector bop_type_; + //! Ternary operator type if applicable, contains a default + //! value at each index corresponding to a ternary op. + std::vector top_type_; + //! Indexes of operands and destination of each instruction. //! The indexes corresponds to positions in the workspace //! where concrete values are hosted. @@ -113,6 +126,10 @@ class NaiveValueMachine { //! each index corresponding to a unary op. std::vector src1_; + //! Operand 2 of each instruction, a default value at + //! each index corresponding to a unary or binary op. + std::vector src2_; + //! Destination of each instruction. std::vector dest_; }; From f58b54914d6c2e76eac0277ea8836e0599fa1041 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Mon, 18 Sep 2023 15:28:07 -0400 Subject: [PATCH 2/5] Add TernaryOps test --- test/test_evaluator.cpp | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/test/test_evaluator.cpp b/test/test_evaluator.cpp index 770dcf66412..30a656593ad 100644 --- a/test/test_evaluator.cpp +++ b/test/test_evaluator.cpp @@ -461,4 +461,41 @@ TEST_F(ExprEvalTest, ReverseArray) { EXPECT_EQ((std::vector)evaluator.evaluate(output), expect); } +//! Test evaluating ternary ops +TEST_F(ExprEvalTest, TernaryOps) { + Fusion fusion; + FusionGuard fg(&fusion); + + ExpressionEvaluator evaluator; + + auto* a = IrBuilder::create(7.0); + auto* b = IrBuilder::create(3.8); + auto* c = IrBuilder::create(0.8); + auto* d = IrBuilder::create(0.2); + auto* t = IrBuilder::create(true); + auto* f = IrBuilder::create(false); + + EXPECT_EQ(evaluator.evaluate(clamp(b, c, a)), b->value()); + EXPECT_EQ(evaluator.evaluate(clamp(a, c, b)), b->value()); + EXPECT_EQ(evaluator.evaluate(clamp(d, c, b)), c->value()); + + EXPECT_EQ( + evaluator.evaluate(lerp(a, b, d)), + a->value() + d->value() * (b->value() - a->value())); + + EXPECT_EQ( + evaluator.evaluate(lerp(a, b, c)), + a->value() + c->value() * (b->value() - a->value())); + EXPECT_EQ( + evaluator.evaluate(lerp(a, b, d)), + a->value() + d->value() * (b->value() - a->value())); + + EXPECT_EQ(evaluator.evaluate(threshold(a, c, b)), a->value()); + EXPECT_EQ(evaluator.evaluate(threshold(d, c, b)), b->value()); + EXPECT_EQ(evaluator.evaluate(threshold(d, d, b)), b->value()); + + EXPECT_EQ(evaluator.evaluate(where(t, a, b)), a->value()); + EXPECT_EQ(evaluator.evaluate(where(f, a, b)), b->value()); +} + } // namespace nvfuser From 3fcaa3677da0a64d56ea49704e661bc8ae1c6c3f Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Mon, 18 Sep 2023 15:28:21 -0400 Subject: [PATCH 3/5] Add ternary ops to evaluator and TernaryOp::evaluate --- csrc/evaluator_common.cpp | 11 +++++++++++ csrc/ir/nodes.cpp | 19 +++++++++++++++---- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/csrc/evaluator_common.cpp b/csrc/evaluator_common.cpp index 58e9314512c..e2ca200f794 100644 --- a/csrc/evaluator_common.cpp +++ b/csrc/evaluator_common.cpp @@ -670,6 +670,17 @@ void NaiveValueMachine::runTernaryOp(int index) { auto& dest = precomputed_values_.values_[dest_index]; switch (top_type_[index]) { + case TernaryOpType::Clamp: + dest = std::min(std::max(a, b), c); + break; + case TernaryOpType::Lerp: + // This is the same lerp computed in helpers.cu + // https://math.stackexchange.com/a/1798323 + dest = (c < 0.5) ? a + c * (b - a) : b - (b - a) * (1.0 - c); + break; + case TernaryOpType::Threshold: + dest = a <= b ? c : a; + break; case TernaryOpType::Where: dest = a ? b : c; break; diff --git a/csrc/ir/nodes.cpp b/csrc/ir/nodes.cpp index f1acc89d9c6..aa8786070c4 100644 --- a/csrc/ir/nodes.cpp +++ b/csrc/ir/nodes.cpp @@ -570,12 +570,23 @@ std::vector TernaryOp::evaluate( const ExpressionEvaluator& ee, const std::vector& inputs) const { using namespace PolymorphicValue_functions; - const auto& in1 = inputs.at(0); - const auto& in2 = inputs.at(1); - const auto& in3 = inputs.at(2); + const auto& a = inputs.at(0); + const auto& b = inputs.at(1); + const auto& c = inputs.at(2); switch (getTernaryOpType()) { + case TernaryOpType::Clamp: + return {std::min(std::max(a, b), c)}; + break; + case TernaryOpType::Lerp: + // This is the same lerp computed in helpers.cu + // https://math.stackexchange.com/a/1798323 + return {(c < 0.5) ? a + c * (b - a) : b - (b - a) * (1.0 - c)}; + break; + case TernaryOpType::Threshold: + return {(a <= b) ? c : a}; + break; case TernaryOpType::Where: - return {in1.as() ? in2 : in3}; + return {a.as() ? b : c}; break; default: NVF_CHECK( From 80cf7ea7a3fbadf02896e14f5493198bde486998 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 19 Sep 2023 08:29:28 -0400 Subject: [PATCH 4/5] Remove changes to LoadStoreOp --- csrc/evaluator_common.cpp | 19 ------------------- csrc/evaluator_common.h | 4 ---- 2 files changed, 23 deletions(-) diff --git a/csrc/evaluator_common.cpp b/csrc/evaluator_common.cpp index e2ca200f794..f71dc030c87 100644 --- a/csrc/evaluator_common.cpp +++ b/csrc/evaluator_common.cpp @@ -355,12 +355,6 @@ NaiveValueMachine::NaiveValueMachine(PrecomputedValues& precomputed_values) makeBinaryOp(bop); } else if (auto top = dynamic_cast(def)) { makeTernaryOp(top); - } else if (auto lsop = dynamic_cast(def)) { - NVF_ERROR( - lsop->opType() == LoadStoreOpType::Set, - "NaiveValueMachine: unsupported LoadStoreOpType: ", - lsop->opType()); - makeSetOp(lsop); } else { // There could be some ops not supported yet. For these ops, we will // bind their outputs. So ignoring them here. @@ -463,19 +457,6 @@ void NaiveValueMachine::makeTernaryOp(TernaryOp* top) { dest_[index] = out; } -void NaiveValueMachine::makeSetOp(LoadStoreOp* lsop) { - int in = lsop->in()->evaluatorIndex(); - int out = lsop->out()->evaluatorIndex(); - - NVF_ERROR(in >= 0, "Integer Machine: unknown input: ", lsop); - NVF_ERROR(out >= 0, "Integer Machine: unknown out: ", lsop); - - int index = makeInstructionEntry(); - inst_type_[index] = InstructionType::SET_OP; - src0_[index] = in; - dest_[index] = out; -} - int NaiveValueMachine::makeInstructionEntry() { int index = num_of_instructions_++; inst_type_.emplace_back(InstructionType::UNARY_OP); diff --git a/csrc/evaluator_common.h b/csrc/evaluator_common.h index 28b4a1e1d9f..18dd88112bd 100644 --- a/csrc/evaluator_common.h +++ b/csrc/evaluator_common.h @@ -58,10 +58,6 @@ class NaiveValueMachine { //! Convert an ternary IR expr to an instruction void makeTernaryOp(TernaryOp* bop); - //! Convert a LoadStoreOp expr to an instruction. This assumes lsop->opType() - //! is equal to LoadStoreOpType::Set. - void makeSetOp(LoadStoreOp* lsop); - //! Create an empty instruction with all default values //! and place it at the end of the instruction buffer. int makeInstructionEntry(); From b2897a6efb4706f660725fa21a6d07f6c4e05712 Mon Sep 17 00:00:00 2001 From: Jacob Hinkle Date: Tue, 19 Sep 2023 09:04:00 -0400 Subject: [PATCH 5/5] Use PrecomputedValues in test --- test/test_evaluator.cpp | 49 +++++++++++++++++++++++------------------ 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/test/test_evaluator.cpp b/test/test_evaluator.cpp index 30a656593ad..f19a2bbb6ff 100644 --- a/test/test_evaluator.cpp +++ b/test/test_evaluator.cpp @@ -475,27 +475,34 @@ TEST_F(ExprEvalTest, TernaryOps) { auto* t = IrBuilder::create(true); auto* f = IrBuilder::create(false); - EXPECT_EQ(evaluator.evaluate(clamp(b, c, a)), b->value()); - EXPECT_EQ(evaluator.evaluate(clamp(a, c, b)), b->value()); - EXPECT_EQ(evaluator.evaluate(clamp(d, c, b)), c->value()); - - EXPECT_EQ( - evaluator.evaluate(lerp(a, b, d)), - a->value() + d->value() * (b->value() - a->value())); - - EXPECT_EQ( - evaluator.evaluate(lerp(a, b, c)), - a->value() + c->value() * (b->value() - a->value())); - EXPECT_EQ( - evaluator.evaluate(lerp(a, b, d)), - a->value() + d->value() * (b->value() - a->value())); - - EXPECT_EQ(evaluator.evaluate(threshold(a, c, b)), a->value()); - EXPECT_EQ(evaluator.evaluate(threshold(d, c, b)), b->value()); - EXPECT_EQ(evaluator.evaluate(threshold(d, d, b)), b->value()); - - EXPECT_EQ(evaluator.evaluate(where(t, a, b)), a->value()); - EXPECT_EQ(evaluator.evaluate(where(f, a, b)), b->value()); + // Run once without PrecomputedValues, then once with + for ([[maybe_unused]] auto i : c10::irange(2)) { + EXPECT_EQ(evaluator.evaluate(clamp(b, c, a)), b->value()); + EXPECT_EQ(evaluator.evaluate(clamp(a, c, b)), b->value()); + EXPECT_EQ(evaluator.evaluate(clamp(d, c, b)), c->value()); + + EXPECT_EQ( + evaluator.evaluate(lerp(a, b, d)), + a->value() + d->value() * (b->value() - a->value())); + + EXPECT_EQ( + evaluator.evaluate(lerp(a, b, c)), + a->value() + c->value() * (b->value() - a->value())); + EXPECT_EQ( + evaluator.evaluate(lerp(a, b, d)), + a->value() + d->value() * (b->value() - a->value())); + + EXPECT_EQ(evaluator.evaluate(threshold(a, c, b)), a->value()); + EXPECT_EQ(evaluator.evaluate(threshold(d, c, b)), b->value()); + EXPECT_EQ(evaluator.evaluate(threshold(d, d, b)), b->value()); + + EXPECT_EQ(evaluator.evaluate(where(t, a, b)), a->value()); + EXPECT_EQ(evaluator.evaluate(where(f, a, b)), b->value()); + + // Now bind a PrecomputedValues + PrecomputedValues pv(&fusion); + evaluator.bindPrecomputedValues(&pv); + } } } // namespace nvfuser