[WIP] New types of DynExpr (comparison, abs, select)

tensorflow/compiler/jit/kernels/xla_ops.cc

@@ -1272,14 +1272,17 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
           VLOG(1) << "input shape is " << ctx->input(input_idx).shape()
                   << ", corresponding xla input shape is " << xla_shape;
           int64_t size = ctx->input(input_idx).shape().dim_size(dim);
-          int64_t dyn_val = expr->solve(size); // TODO: check if the result is correct later.
+          std::optional<int64_t> dyn_val =
+              expr->solve(size);  // TODO: check if the result is correct later.
           VLOG(1) << "Found dynamic input. Real size is: " << size
-                        << ", solved dynamic value is " << dyn_val;
-          if (dyn_val == -1) {
+                  << ", solved dynamic value is "
+                  << (dyn_val.has_value() ? std::to_string(*dyn_val)
+                                          : std::string("<none>"));
+          if (!dyn_val.has_value()) {
             VLOG(1) << "Warning: Failed to solve the expression";
             continue;
           }
-          dyn_vals.insert(dyn_val);
+          dyn_vals.insert(*dyn_val);
         }
       }
     }

tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc

@@ -313,8 +313,8 @@ class StridedSliceOp : public XlaOpKernel {
           slice_begin.push_back(begin[i]);
           slice_begin_expr.push_back(begin_expr[i]);
           slice_end.push_back(std::max(end[i], begin[i]));
-          slice_end_expr.push_back((end[i] > begin[i]) ? end_expr[i]
-                                                       : begin_expr[i]);
+          slice_end_expr.push_back(
+              xla::dexpr::Max(*end_expr[i], *begin_expr[i])->s());
           slice_strides.push_back(strides[i]);
         } else {
           // Negative stride: swap begin and end, add 1 because the interval
@@ -326,9 +326,9 @@ class StridedSliceOp : public XlaOpKernel {
           slice_end.push_back(std::max(input_shape.dim_size(i) - end[i] - 1,
                                        input_shape.dim_size(i) - begin[i] - 1));
           slice_end_expr.push_back(
-              (end[i] < begin[i])
-                  ? (*input_exprs[i] - *end_expr[i] - 1)->s()
-                  : (*input_exprs[i] - *begin_expr[i] - 1)->s());
+              xla::dexpr::Max(*(*input_exprs[i] - *end_expr[i] - 1)->s(),
+                              *(*input_exprs[i] - *begin_expr[i] - 1)->s())
+                  ->s());
           slice_strides.push_back(-strides[i]);
           dimensions_to_reverse.push_back(i);
         }

tensorflow/core/util/strided_slice_op.cc

@@ -371,21 +371,26 @@ absl::Status ValidateStridedSliceOp(
                                         : x_fwd;
       }
     };
-    auto canonical_expr = [stride_i, dim_i, masks, valid_range,
-                           valid_range_expr, dim_i_expr](int64_t x, int c) {
+    auto canonical_expr = [stride_i, masks, valid_range_expr,
+                           dim_i_expr](xla::DynExpr* x_expr, int c) {
       if (masks[c]) {
         return stride_i > 0 ? valid_range_expr[c]
                             : valid_range_expr[(c + 1) & 1];
       } else {
-        int64_t x_fwd =
-            x < 0 ? dim_i + x : x;  // make negative indices positive
-        xla::DynExpr* x_expr = xla::DynExpr::_(x);
+        xla::DynExpr* wrapped_x_expr = (*dim_i_expr + *x_expr)->s();
         xla::DynExpr* x_fwd_expr =
-            x < 0 ? (*dim_i_expr + *x_expr)
-                  : x_expr;  // make negative indices positive
-        return x_fwd < valid_range[0]   ? valid_range_expr[0]
-               : x_fwd > valid_range[1] ? valid_range_expr[1]
-                                        : x_fwd_expr;
+            xla::dexpr::Select(*xla::dexpr::Gt(0, *x_expr), *wrapped_x_expr,
+                               *x_expr)
+                ->s();
+        xla::DynExpr* low_clamped_expr =
+            xla::dexpr::Select(
+                *xla::dexpr::Gt(*valid_range_expr[0], *x_fwd_expr),
+                *valid_range_expr[0], *x_fwd_expr)
+                ->s();
+        return xla::dexpr::Select(
+                   *xla::dexpr::Gt(*low_clamped_expr, *valid_range_expr[1]),
+                   *valid_range_expr[1], *low_clamped_expr)
+            ->s();
       }
     };
     if (shrink_i && stride_i <= 0) {
@@ -403,9 +408,13 @@ absl::Status ValidateStridedSliceOp(
         // and canonical puts these to n-1 and 0, which implies a degenerate
         // interval. Fortunately, it is now safe to re-create end as begin+1.
         int64_t x_fwd = begin_i < 0 ? dim_i + begin_i : begin_i;
-        xla::DynExpr* x_fwd_expr = begin_i < 0
-                                       ? (*dim_i_expr + *(*begin_expr)[i])->s()
-                                       : (*begin_expr)[i];
+        xla::DynExpr* wrapped_x_expr =
+            (*dim_i_expr + *(*begin_expr)[i])->s();
+        xla::DynExpr* x_fwd_expr =
+            xla::dexpr::Select(
+                *xla::dexpr::Gt(0, *(*begin_expr)[i]), *wrapped_x_expr,
+                *(*begin_expr)[i])
+                ->s();
         begin_i = x_fwd;
         end_i = begin_i + 1;
 
@@ -422,10 +431,10 @@ absl::Status ValidateStridedSliceOp(
         begin_i = canonical(begin_raw, 0);
         end_i = canonical(end_raw, 1);
         if (begin_expr) {
-          (*begin_expr)[i] = canonical_expr(begin_raw, 0)->s();
+          (*begin_expr)[i] = canonical_expr((*begin_expr)[i], 0)->s();
         }
         if (end_expr) {
-          (*end_expr)[i] = canonical_expr(end_raw, 1)->s();
+          (*end_expr)[i] = canonical_expr((*end_expr)[i], 1)->s();
         }
       }
       // Update optimization values

third_party/xla/xla/service/llvm_ir/llvm_util.cc

@@ -897,10 +897,12 @@ llvm::Value* EmitExpression(llvm::IRBuilderBase* b, DynExpr* expr) {
     llvm::Value* v_rhs = EmitExpression(b, mul_node->get_rhs());
     return b->CreateMul(v_lhs, v_rhs, "mul_dims");
   }
-  // TODO: Check if this should ever happen
   if (Div* div_node = dynamic_cast<Div*>(expr)) {
     llvm::Value* v_lhs = EmitExpression(b, div_node->get_lhs());
     llvm::Value* v_rhs = EmitExpression(b, div_node->get_rhs());
+    if (auto* rhs_const = llvm::dyn_cast<llvm::ConstantInt>(v_rhs)) {
+      CHECK_NE(rhs_const->getSExtValue(), 0);
+    }
     return b->CreateUDiv(v_lhs, v_rhs, "div_dims");
   }
   if (Add* add_node = dynamic_cast<Add*>(expr)) {
@@ -913,6 +915,47 @@ llvm::Value* EmitExpression(llvm::IRBuilderBase* b, DynExpr* expr) {
     llvm::Value* v_rhs = EmitExpression(b, sub_node->get_rhs());
     return b->CreateSub(v_lhs, v_rhs, "sub_dims");
   }
+  if (GtExpr* gt_node = dynamic_cast<GtExpr*>(expr)) {
+    llvm::Value* v_lhs = EmitExpression(b, gt_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpression(b, gt_node->get_rhs());
+    return b->CreateICmpSGT(v_lhs, v_rhs, "gt_dims");
+  }
+  if (EqExpr* eq_node = dynamic_cast<EqExpr*>(expr)) {
+    llvm::Value* v_lhs = EmitExpression(b, eq_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpression(b, eq_node->get_rhs());
+    return b->CreateICmpEQ(v_lhs, v_rhs, "eq_dims");
+  }
+  if (NeExpr* ne_node = dynamic_cast<NeExpr*>(expr)) {
+    llvm::Value* v_lhs = EmitExpression(b, ne_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpression(b, ne_node->get_rhs());
+    return b->CreateICmpNE(v_lhs, v_rhs, "ne_dims");
+  }
+  if (GeExpr* ge_node = dynamic_cast<GeExpr*>(expr)) {
+    llvm::Value* v_lhs = EmitExpression(b, ge_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpression(b, ge_node->get_rhs());
+    return b->CreateICmpSGE(v_lhs, v_rhs, "ge_dims");
+  }
+  if (MaxExpr* max_node = dynamic_cast<MaxExpr*>(expr)) {
+    llvm::Value* v_lhs = EmitExpression(b, max_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpression(b, max_node->get_rhs());
+    llvm::Value* pred = b->CreateICmpSGT(v_lhs, v_rhs, "max_dims_pred");
+    return b->CreateSelect(pred, v_lhs, v_rhs, "max_dims");
+  }
+  if (AbsExpr* abs_node = dynamic_cast<AbsExpr*>(expr)) {
+    llvm::Value* value = EmitExpression(b, abs_node->get_expr());
+    llvm::Value* is_negative =
+        b->CreateICmpSLT(value, llvm::ConstantInt::get(i64Type, 0, true),
+                         "abs_dims_neg");
+    llvm::Value* negated = b->CreateNeg(value, "abs_dims_negated");
+    return b->CreateSelect(is_negative, negated, value, "abs_dims");
+  }
+  if (SelectExpr* select_node = dynamic_cast<SelectExpr*>(expr)) {
+    llvm::Value* pred = EmitExpression(b, select_node->get_pred());
+    llvm::Value* on_true = EmitExpression(b, select_node->get_on_true());
+    llvm::Value* on_false = EmitExpression(b, select_node->get_on_false());
+    CHECK(pred->getType()->isIntegerTy(1));
+    return b->CreateSelect(pred, on_true, on_false, "select_dims");
+  }
   return nullptr;
 }