Add support for IfThenElse (pytorch#103)

Author: resistor

test/cpp/tensorexpr/test_expr.cpp

@@ -324,5 +324,31 @@ void testCond01() {
   ExpectAllNear(a_v, a_ref, 1e-5);
 }
 
+void testIfThenElse01() {
+  KernelScope kernel_scope;
+  Expr v = ifThenElse(Expr(1), Expr(1.0f), Expr(2.0f));
+
+  std::ostringstream oss;
+  oss << v;
+  ASSERT_EQ(oss.str(), "IfThenElse(1, 1, 2)");
+
+  SimpleIREvaluator eval(v);
+  eval();
+  ASSERT_EQ(eval.value().as<float>(), 1.0f);
+}
+
+void testIfThenElse02() {
+  KernelScope kernel_scope;
+  Expr v = ifThenElse(Expr(0), Expr(1.0f), Expr(2.0f));
+
+  std::ostringstream oss;
+  oss << v;
+  ASSERT_EQ(oss.str(), "IfThenElse(0, 1, 2)");
+
+  SimpleIREvaluator eval(v);
+  eval();
+  ASSERT_EQ(eval.value().as<float>(), 2.0f);
+}
+
 } // namespace jit
 } // namespace torch

test/cpp/tensorexpr/test_llvm.cpp

@@ -153,6 +153,30 @@ void testLLVMLoadStoreTest() {
   EXPECT_EQ(b_buffer[0], 42);
 }
 
+void testLLVMIfThenElseTest() {
+  KernelScope kernel_scope;
+  Buffer a(Var("A", kHandle), kInt32, {1});
+  Buffer b(Var("B", kHandle), kInt32, {1});
+  Buffer c(Var("C", kHandle), kInt32, {1});
+  std::vector<int32_t> a_buffer = {42};
+  std::vector<int32_t> b_buffer = {-11};
+  std::vector<int32_t> c_buffer = {1};
+
+  auto store = Store::make(
+      b,
+      IntImm::make(0),
+      IfThenElse::make(
+        Load::make(c, IntImm::make(0), IntImm::make(1)), // cond
+        Load::make(a, IntImm::make(0), IntImm::make(1)), // then
+        IntImm::make(0)), // else
+      IntImm::make(1));
+  LLVMCodeGen cg(store, {a, b, c});
+  std::vector<void*> args({a_buffer.data(), b_buffer.data(), c_buffer.data()});
+  EXPECT_EQ(cg.value<int>(args), 0);
+  EXPECT_EQ(a_buffer[0], 42);
+  EXPECT_EQ(b_buffer[0], 42);
+}
+
 void testLLVMVecLoadStoreTest() {
   KernelScope kernel_scope;
   Buffer a(Var("A", kHandle), kInt32, {1});

test/cpp/tensorexpr/tests.h

@@ -41,6 +41,8 @@ namespace jit {
   _(AsmjitIntMulTest)           \
   _(AsmjitIntDivTest)           \
   _(Cond01)                     \
+  _(IfThenElse01)               \
+  _(IfThenElse02)               \
   _(ATen_cast_Float)            \
   _(ATennegInt)                 \
   _(ATennegFloat)               \
@@ -109,7 +111,8 @@ namespace jit {
   _(LLVMComputeMul)             \
   _(LLVMBroadcastAdd)           \
   _(LLVMDynamicShapeAdd)        \
-  _(LLVMBindDynamicShapeAdd)
+  _(LLVMBindDynamicShapeAdd)    \
+  _(LLVMIfThenElseTest)
 
 #define TH_FORALL_TESTS_CUDA(_) \
   _(CudaTestVectorAdd01)        \

torch/csrc/jit/tensorexpr/eval.h

@@ -379,6 +379,15 @@ class SimpleIREvaluator : public CodeGen, public IRVisitor {
     }
   }
 
+  TORCH_API void visit(const IfThenElse* v) override {
+    v->condition().accept(this);
+    if (value_.as<int>()) {
+      v->true_value().accept(this);
+    } else {
+      v->false_value().accept(this);
+    }
+  }
+
   TORCH_API void visit(const Load* v) override {
     const Variable* base_node = v->base_handle().node();
     auto iter = buffer_mapping_.find(base_node);

torch/csrc/jit/tensorexpr/expr.cpp

@@ -193,6 +193,11 @@ Expr fmod(const Expr& v1, const Expr& v2) {
 Expr remainder(const Expr& v1, const Expr& v2) {
   return Intrinsics::make(kRemainder, v1, v2);
 }
+
+Expr ifThenElse(const Expr& c, const Expr& t, const Expr& f) {
+  return IfThenElse::make(c, t, f);
+}
+
 } // namespace tensorexpr
 } // namespace jit
 } // namespace torch

torch/csrc/jit/tensorexpr/expr.h

@@ -261,6 +261,9 @@ TORCH_API Expr pow(const Expr& v1, const Expr& v2);
 TORCH_API Expr fmod(const Expr& v1, const Expr& v2);
 TORCH_API Expr remainder(const Expr& v1, const Expr& v2);
 
+TORCH_API Expr ifThenElse(const Expr& c, const Expr& t, const Expr& f);
+
+
 } // namespace tensorexpr
 } // namespace jit
 } // namespace torch

torch/csrc/jit/tensorexpr/ir.h

@@ -606,6 +606,40 @@ class Broadcast : public ExprNode<Broadcast> {
   Expr value_;
   int lanes_;
 };
+class IfThenElse : public ExprNode<IfThenElse> {
+ public:
+  const Expr& condition() const {
+    return condition_;
+  }
+
+  // Lazily evaluated only if condition is true
+  const Expr& true_value() const {
+    return true_;
+  }
+
+  // Lazily evaluated only if condition is false
+  const Expr& false_value() const {
+    return false_;
+  }
+
+  static Expr make(const Expr& c, const Expr& t, const Expr& f) {
+    return Expr(new IfThenElse(c, t, f));
+  }
+
+ private:
+  IfThenElse(const Expr& c, const Expr& t, const Expr& f)
+      : ExprNodeBase(t.dtype()),
+        condition_(c),
+        true_(t),
+        false_(f) {
+    CHECK_EQ(c.dtype().scalar_type(), kInt32);
+    CHECK_EQ(c.dtype().lanes(), 1);
+    CHECK_EQ(t.dtype(), f.dtype());
+  }
+  Expr condition_;
+  Expr true_;
+  Expr false_;
+};
 
 class BaseCallNode : public BaseExprNode {
  public:

torch/csrc/jit/tensorexpr/ir_mutator.cpp

@@ -150,6 +150,22 @@ Expr IRMutator::mutate(const Broadcast* v) {
   return Broadcast::make(value_new, lanes);
 }
 
+Expr IRMutator::mutate(const IfThenElse* v) {
+  Expr condition = v->condition();
+  Expr true_value = v->true_value();
+  Expr false_value = v->false_value();
+  Expr condition_new = condition.accept_mutator(this);
+  Expr true_value_new = true_value.accept_mutator(this);
+  Expr false_value_new = false_value.accept_mutator(this);
+  if (same_node(condition, condition_new) &&
+      same_node(true_value, true_value_new) &&
+      same_node(false_value, false_value_new)) {
+    return Expr(v);
+  }
+
+  return IfThenElse::make(condition_new, true_value_new, false_value_new);
+}
+
 Expr IRMutator::mutate(const Intrinsics* v) {
   const BaseCallNode* base = v;
   return this->mutate(base);

torch/csrc/jit/tensorexpr/ir_mutator.h

@@ -24,6 +24,7 @@ class For;
 class Block;
 class Store;
 class Broadcast;
+class IfThenElse;
 class Expr;
 class Stmt;
 class BaseCallNode;
@@ -52,6 +53,7 @@ class TORCH_API IRMutator {
   virtual Expr mutate(const Ramp* v);
   virtual Expr mutate(const Load* v);
   virtual Expr mutate(const Broadcast* v);
+  virtual Expr mutate(const IfThenElse* v);
   // BaseCallNode is the base class for all call nodes.
   // For any visitors that only needs the common behavior, only override this
   // function is enough. This is because all derived class handlers will call

torch/csrc/jit/tensorexpr/ir_printer.cpp

@@ -161,6 +161,10 @@ void IRPrinter::visit(const Broadcast* v) {
   os() << "Broadcast(" << v->value() << ", " << v->lanes() << ")";
 }
 
+void IRPrinter::visit(const IfThenElse* v) {
+  os() << "IfThenElse(" << v->condition() << ", " << v->true_value() << ", " << v->false_value() << ")";
+}
+
 void IRPrinter::visit(const BaseCallNode* v) {
   os() << v->func_name() << "(";
   for (int i = 0; i < v->nparams(); i++) {