Add support for pext/pdep (#1320)

Author: dtcxzyw

ir/instr.cpp

@@ -187,6 +187,8 @@ void BinOp::print(ostream &os) const {
   case UCmp:          str = "ucmp "; break;
   case SCmp:          str = "scmp "; break;
   case Clmul:         str = "clmul "; break;
+  case PExt:          str = "pext "; break;
+  case PDep:          str = "pdep "; break;
   }
 
   os << getName() << " = " << str;
@@ -479,6 +481,16 @@ StateValue BinOp::toSMT(State &s) const {
       return {a.clmul(b), ap && bp};
     };
     break;
+  case PExt:
+    fn = [&](auto &a, auto &ap, auto &b, auto &bp) -> StateValue {
+      return {a.pext(b), ap && bp};
+    };
+    break;
+  case PDep:
+    fn = [&](auto &a, auto &ap, auto &b, auto &bp) -> StateValue {
+      return {a.pdep(b), ap && bp};
+    };
+    break;
   }
 
   function<pair<StateValue,StateValue>(const expr&, const expr&, const expr&,

ir/instr.h

@@ -41,7 +41,7 @@ class BinOp final : public Instr {
             SAdd_Overflow, UAdd_Overflow, SSub_Overflow, USub_Overflow,
             SMul_Overflow, UMul_Overflow,
             And, Or, Xor, Cttz, Ctlz, UMin, UMax, SMin, SMax, Abs,
-            UCmp, SCmp, Clmul };
+            UCmp, SCmp, Clmul, PExt, PDep };
   enum Flags { None = 0, NSW = 1 << 0, NUW = 1 << 1, Exact = 1 << 2, Disjoint = 1 << 3 };
 
 private:

llvm_util/llvm2alive.cpp

@@ -841,7 +841,9 @@ class llvm2alive_ : public llvm::InstVisitor<llvm2alive_, unique_ptr<Instr>> {
     case llvm::Intrinsic::abs:
     case llvm::Intrinsic::ucmp:
     case llvm::Intrinsic::scmp:
-    case llvm::Intrinsic::clmul: {
+    case llvm::Intrinsic::clmul:
+    case llvm::Intrinsic::pext:
+    case llvm::Intrinsic::pdep: {
       PARSE_BINOP();
       addNoundefAssumes(i, {a, b});
       BinOp::Op op;
@@ -868,6 +870,8 @@ class llvm2alive_ : public llvm::InstVisitor<llvm2alive_, unique_ptr<Instr>> {
       case llvm::Intrinsic::ucmp:     op = BinOp::UCmp; break;
       case llvm::Intrinsic::scmp:     op = BinOp::SCmp; break;
       case llvm::Intrinsic::clmul:    op = BinOp::Clmul; break;
+      case llvm::Intrinsic::pext:     op = BinOp::PExt; break;
+      case llvm::Intrinsic::pdep:     op = BinOp::PDep; break;
       default: UNREACHABLE();
       }
       ret = make_unique<BinOp>(*ty, value_name(i), *a, *b, op);

smt/expr.cpp

@@ -1136,6 +1136,36 @@ expr expr::clmul(const expr &rhs) const {
   return res;
 }
 
+expr expr::pext(const expr &mask) const {
+  C(mask);
+  auto nbits = bits();
+
+  auto res = mkUInt(0, sort());
+  auto shamt = mkUInt(0, sort());
+  for (unsigned i = 0; i < nbits; ++i) {
+    auto active = mask.extract(i, i).zext(nbits - 1);
+    res = res | ((active & extract(i, i).zext(nbits - 1)) << shamt);
+    shamt = shamt + active;
+  }
+
+  return res;
+}
+
+expr expr::pdep(const expr &mask) const {
+  C(mask);
+  auto nbits = bits();
+
+  auto res = mkUInt(0, sort());
+  auto shamt = mkUInt(0, sort());
+  for (unsigned i = 0; i < nbits; ++i) {
+    auto active = mask.extract(i, i).zext(nbits - 1);
+    res = res | ((active & lshr(shamt)) << mkUInt(i, sort()));
+    shamt = shamt + active;
+  }
+
+  return res;
+}
+
 expr expr::umin(const expr &rhs) const {
   return mkIf(ule(rhs), *this, rhs);
 }

smt/expr.h

@@ -218,6 +218,8 @@ class expr {
   expr ctlz() const;
   expr ctpop() const;
   expr clmul(const expr &rhs) const;
+  expr pext(const expr &mask) const;
+  expr pdep(const expr &mask) const;
 
   expr umin(const expr &rhs) const;
   expr umax(const expr &rhs) const;

tests/alive-tv/pext-pdep.srctgt.ll

@@ -0,0 +1,10 @@
+define i8 @src(i8 %val, i8 %mask) {
+  %pext = call i8 @llvm.pext.i8(i8 %val, i8 %mask)
+  %pdep = call i8 @llvm.pdep.i8(i8 %pext, i8 %mask)
+  ret i8 %pdep
+}
+
+define i8 @tgt(i8 %val, i8 %mask) {
+  %masked = and i8 %val, %mask
+  ret i8 %masked
+}

tests/unit/pdep.opt

@@ -0,0 +1,19 @@
+Name: pdep constant1
+%r = pdep i8 10, i8 204
+  =>
+%r = 136
+
+Name: pdep constant2
+%r = pdep i8 15, i8 170
+  =>
+%r = 170
+
+Name: pdep simplify1
+%r = pdep i8 %x, i8 0
+  =>
+%r = 0
+
+Name: pdep simplify2
+%r = pdep i8 %x, i8 255
+  =>
+%r = %x

tests/unit/pext.opt

@@ -0,0 +1,19 @@
+Name: pext constant1
+%r = pext i8 170, i8 204
+  =>
+%r = 10
+
+Name: pext constant2
+%r = pext i8 255, i8 170
+  =>
+%r = 15
+
+Name: pext simplify1
+%r = pext i8 %x, i8 0
+  =>
+%r = 0
+
+Name: pext simplify2
+%r = pext i8 %x, i8 255
+  =>
+%r = %x

tools/alive_lexer.re

@@ -287,6 +287,8 @@ space+ {
 "ucmp" { return UCMP; }
 "scmp" { return SCMP; }
 "clmul" { return CLMUL; }
+"pext" { return PEXT; }
+"pdep" { return PDEP; }
 "oeq" { return OEQ; }
 "ogt" { return OGT; }
 "oge" { return OGE; }

tools/alive_parser.cpp

@@ -711,6 +711,8 @@ static unsigned parse_binop_flags(token op_token) {
   case UCMP:
   case SCMP:
   case CLMUL:
+  case PEXT:
+  case PDEP:
     return BinOp::None;
   default:
     UNREACHABLE();
@@ -784,6 +786,8 @@ static unique_ptr<Instr> parse_binop(string_view name, token op_token) {
   case UCMP: op = BinOp::UCmp; break;
   case SCMP: op = BinOp::SCmp; break;
   case CLMUL: op = BinOp::Clmul; break;
+  case PEXT: op = BinOp::PExt; break;
+  case PDEP: op = BinOp::PDep; break;
   default:
     UNREACHABLE();
   }
@@ -1276,6 +1280,8 @@ static unique_ptr<Instr> parse_instr(string_view name) {
   case UCMP:
   case SCMP:
   case CLMUL:
+  case PEXT:
+  case PDEP:
     return parse_binop(name, t);
   case FADD:
   case FSUB: