Add support for scalable vectors

ir/function.cpp

@@ -3,6 +3,8 @@
 
 #include "ir/function.h"
 #include "ir/instr.h"
+#include "ir/type.h"
+#include "util/config.h"
 #include "util/errors.h"
 #include "util/hash.h"
 #include "util/sort.h"
@@ -168,9 +170,14 @@ expr Function::getTypeConstraints() const {
       t &= v.getTypeConstraints();
     }
   }
+
   return t;
 }
 
+expr Function::getVScaleExpr() const {
+  return VectorType::getVScaleVar();
+}
+
 void Function::rauw(const Value &what, Value &with) {
   for (auto bb : getBBs())
     bb->rauw(what, with);

ir/function.h

@@ -183,6 +183,7 @@ class Function final {
   bool hasReturn() const;
   unsigned bitsPointers() const { return bits_pointers; }
   unsigned bitsPtrOffset() const { return bits_ptr_offset; }
+  smt::expr getVScaleExpr() const;
   bool isLittleEndian() const { return little_endian; }
   bool isVarArgs() const { return is_var_args; }
 

ir/instr.cpp

@@ -974,6 +974,35 @@ unique_ptr<Instr> FpBinOp::dup(Function &f, const string &suffix) const {
 }
 
 
+vector<Value*> VScale::operands() const {
+  return {};
+}
+
+bool VScale::propagatesPoison() const { return false; }
+
+bool VScale::hasSideEffects() const { return false; }
+
+void VScale::rauw(const Value &what, Value &with) {}
+
+void VScale::print(ostream &os) const {
+  os << getName() << " = vscale";
+}
+
+StateValue VScale::toSMT(State &s) const {
+  auto e = s.getFn().getVScaleExpr();
+  return { e.zextOrTrunc(bits()), true };
+}
+
+expr VScale::getTypeConstraints(const Function &f) const {
+  return Value::getTypeConstraints() &&
+         getType().enforceIntType();
+}
+
+unique_ptr<Instr> VScale::dup(Function &f, const string &suffix) const {
+  return make_unique<VScale>(getType(), getName() + suffix);
+}
+
+
 vector<Value*> UnaryOp::operands() const {
   return { val };
 }

ir/instr.h

@@ -102,6 +102,23 @@ class FpBinOp final : public Instr {
 };
 
 
+class VScale final : public Instr {
+public:
+  VScale(Type &type, std::string &&name)
+    : Instr(type, std::move(name)) {}
+
+  std::vector<Value*> operands() const override;
+  bool propagatesPoison() const override;
+  bool hasSideEffects() const override;
+  void rauw(const Value &what, Value &with) override;
+  void print(std::ostream &os) const override;
+  StateValue toSMT(State &s) const override;
+  smt::expr getTypeConstraints(const Function &f) const override;
+  std::unique_ptr<Instr>
+    dup(Function &f, const std::string &suffix) const override;
+};
+
+
 class UnaryOp final : public Instr {
 public:
   enum Op {

ir/type.cpp

@@ -6,6 +6,7 @@
 #include "ir/state.h"
 #include "smt/solver.h"
 #include "util/compiler.h"
+#include "util/config.h"
 #include <array>
 #include <cassert>
 #include <numeric>
@@ -1089,12 +1090,26 @@ void ArrayType::print(ostream &os) const {
 }
 
 
-VectorType::VectorType(string &&name, unsigned elements, Type &elementTy)
-  : AggregateType(std::move(name), false) {
-  assert(elements != 0);
-  this->elements = elements;
+expr VectorType::getVScaleVar() {
+  return expr::mkVar("vscale", var_vector_elements);
+}
+
+expr VectorType::vscale() const {
+  if (!scalable)
+    return expr::mkUInt(1, var_vector_elements);
+  return defined ? expr::mkUInt(vscale_value, var_vector_elements)
+                 : getVScaleVar();
+}
+
+VectorType::VectorType(string &&name, unsigned elems, Type &elTy, bool scal)
+  : AggregateType(std::move(name), false), scalable(scal), min_elements(elems),
+    vscale_value(scal ? util::config::vscale_value : 1) {
+  assert(elems != 0);
+  if (scalable)
+    elems *= vscale_value;
+  this->elements = elems;
   defined = true;
-  children.resize(elements, &elementTy);
+  children.resize(elements, &elTy);
   is_padding.resize(elements, false);
 }
 
@@ -1157,6 +1172,12 @@ expr VectorType::getTypeConstraints() const {
     r &= numElements().ugt(i).implies(elementTy == *children[i]);
   }
 
+  if (scalable && defined) {
+    r &= vscale() == expr::mkUInt(vscale_value, var_vector_elements);
+  } else if (scalable) {
+    r &= vscale().uge(expr::mkUInt(1, var_vector_elements));
+  }
+
   return r;
 }
 
@@ -1172,14 +1193,25 @@ bool VectorType::isVectorType() const {
   return true;
 }
 
+void VectorType::fixup(const Model &m) {
+  if (!defined && scalable)
+    vscale_value = m.getUInt(vscale());
+  AggregateType::fixup(m);
+}
+
 expr VectorType::enforceVectorType(
     const function<expr(const Type&)> &enforceElem) const {
   return enforceElem(*children[0]);
 }
 
 void VectorType::print(ostream &os) const {
-  if (elements)
-    os << '<' << elements << " x " << *children[0] << '>';
+  if (!elements)
+    return;
+  os << '<';
+  if (scalable) {
+    os << "vscale" << ":" << vscale_value << " x ";
+  }
+  os << min_elements << " x " << *children[0] << '>';
 }
 
 

ir/type.h

@@ -5,6 +5,7 @@
 
 #include "ir/attrs.h"
 #include "smt/expr.h"
+#include "util/config.h"
 
 #include <functional>
 #include <memory>
@@ -335,9 +336,19 @@ class ArrayType final : public AggregateType {
 
 
 class VectorType final : public AggregateType {
+private:
+  bool scalable = false;
+  unsigned min_elements, vscale_value;
+
 public:
+  smt::expr vscale() const;
+  static smt::expr getVScaleVar();
+
+  bool isScalable() const { return scalable; }
+  unsigned getMinElements() const { return min_elements; }
+
   VectorType(std::string &&name) : AggregateType(std::move(name)) {}
-  VectorType(std::string &&name, unsigned elements, Type &elementTy);
+  VectorType(std::string &&name, unsigned elems, Type &elemTy, bool scalable);
 
   IR::StateValue extract(const IR::StateValue &vector,
                          const smt::expr &index) const;
@@ -348,6 +359,7 @@ class VectorType final : public AggregateType {
   unsigned maxSubBitAccess() const override;
   smt::expr scalarSize() const override;
   bool isVectorType() const override;
+  void fixup(const smt::Model &m) override;
   smt::expr enforceVectorType(
     const std::function<smt::expr(const Type&)> &enforceElem) const override;
   void print(std::ostream &os) const override;

llvm_util/cmd_args_def.h

@@ -24,6 +24,7 @@ config::debug = opt_debug;
 config::quiet = opt_quiet;
 config::max_offset_bits = opt_max_offset_in_bits;
 config::max_sizet_bits  = opt_max_sizet_in_bits;
+config::vscale_value = opt_vscale;
 
 if ((config::disallow_ub_exploitation = opt_disallow_ub_exploitation)) {
   config::disable_undef_input = true;

llvm_util/cmd_args_list.h

@@ -185,4 +185,9 @@ llvm::cl::opt<bool> opt_disallow_ub_exploitation(
   llvm::cl::desc("Disallow UB exploitation by optimizations (default=allow)"),
   llvm::cl::init(false), llvm::cl::cat(alive_cmdargs));
 
+llvm::cl::opt<unsigned> opt_vscale(LLVM_ARGS_PREFIX "vscale",
+  llvm::cl::desc("Set vscale value for scalable vectors (default=2)"),
+  llvm::cl::init(2), llvm::cl::value_desc("value"),
+  llvm::cl::cat(alive_cmdargs));
+
 }

llvm_util/llvm2alive.cpp

@@ -2,6 +2,7 @@
 // Distributed under the MIT license that can be found in the LICENSE file.
 
 #include "llvm_util/llvm2alive.h"
+#include "ir/type.h"
 #include "ir/x86_intrinsics.h"
 #include "llvm_util/known_fns.h"
 #include "llvm_util/utils.h"
@@ -1224,6 +1225,12 @@ class llvm2alive_ : public llvm::InstVisitor<llvm2alive_, unique_ptr<Instr>> {
       addNoundefAssumes(i, {a, b});
       return make_unique<VaCopy>(*a, *b);
     }
+    case llvm::Intrinsic::vscale: {
+      auto ty = llvm_type2alive(i.getType());
+      if (!ty)
+        return error(i);
+      return make_unique<VScale>(*ty, value_name(i));
+    }
 
     // do nothing intrinsics
     case llvm::Intrinsic::dbg_declare:
@@ -1335,8 +1342,17 @@ class llvm2alive_ : public llvm::InstVisitor<llvm2alive_, unique_ptr<Instr>> {
   RetTy visitShuffleVectorInst(llvm::ShuffleVectorInst &i) {
     PARSE_BINOP();
     vector<unsigned> mask;
-    for (auto m : i.getShuffleMask())
-      mask.push_back(m);
+
+    unsigned replicate = 1;
+    if (i.getType()->isScalableTy()) {
+      replicate = config::vscale_value;
+    }
+
+    auto &&sm = i.getShuffleMask();
+    for (unsigned j = 0; j < replicate; j++) {
+      mask.insert(mask.end(), sm.begin(), sm.end());
+    }
+
     return
       make_unique<ShuffleVector>(*ty, value_name(i), *a, *b, std::move(mask));
   }

llvm_util/utils.cpp

@@ -203,7 +203,8 @@ Type* llvm_type2alive(const llvm::Type *ty) {
     return cache.get();
   }
   // TODO: non-fixed sized vectors
-  case llvm::Type::FixedVectorTyID: {
+  case llvm::Type::FixedVectorTyID:
+  case llvm::Type::ScalableVectorTyID: {
     auto &cache = type_cache[ty];
     if (!cache) {
       auto vty = cast<llvm::VectorType>(ty);
@@ -212,7 +213,7 @@ Type* llvm_type2alive(const llvm::Type *ty) {
       if (!ety || elems > 1024)
         return nullptr;
       cache = make_unique<VectorType>("ty_" + to_string(type_id_counter++),
-                                      elems, *ety);
+                                      elems, *ety, vty->isScalableTy());
     }
     return cache.get();
   }

tests/alive-tv/vector/scalable/insert-loop1.ll

@@ -0,0 +1,29 @@
+define i32 @src() {
+  ret i32 1
+}
+
+define i32 @tgt() {
+entry:
+  %vs = call i32 @llvm.vscale.i32()
+  %len = mul i32 %vs, 2
+  %init_vec = insertelement <vscale x 2 x i32> poison, i32 0, i32 0
+  br label %for.cond
+
+for.cond:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %vec = phi <vscale x 2 x i32> [ %init_vec, %entry ], [ %new_vec, %for.body ]
+  %cmp = icmp slt i32 %i, %len
+  br i1 %cmp, label %for.body, label %exit
+
+for.body:
+  %new_vec = insertelement <vscale x 2 x i32> %vec, i32 %i, i32 %i
+  %inc = add i32 %i, 1
+  br label %for.cond
+
+exit:
+  %last_idx = sub i32 %len, 1
+  %result = extractelement <vscale x 2 x i32> %vec, i32 %last_idx
+  ret i32 %result
+}
+
+; TEST-ARGS: --vscale=1 -tgt-unroll=2

tests/alive-tv/vector/scalable/insert-loop2.ll

@@ -0,0 +1,28 @@
+define i32 @src() {
+  ret i32 3
+}
+define i32 @tgt() {
+entry:
+  %vs = call i32 @llvm.vscale.i32()
+  %len = mul i32 %vs, 2
+  %init_vec = insertelement <vscale x 2 x i32> poison, i32 0, i32 0
+  br label %for.cond
+
+for.cond:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %vec = phi <vscale x 2 x i32> [ %init_vec, %entry ], [ %new_vec, %for.body ]
+  %cmp = icmp slt i32 %i, %len
+  br i1 %cmp, label %for.body, label %exit
+
+for.body:
+  %new_vec = insertelement <vscale x 2 x i32> %vec, i32 %i, i32 %i
+  %inc = add i32 %i, 1
+  br label %for.cond
+
+exit:
+  %last_idx = sub i32 %len, 1
+  %result = extractelement <vscale x 2 x i32> %vec, i32 %last_idx
+  ret i32 %result
+}
+
+; TEST-ARGS: --vscale=2 -tgt-unroll=4

tests/alive-tv/vector/scalable/insert-loop3.ll

@@ -0,0 +1,28 @@
+define i32 @src() {
+  ret i32 5
+}
+define i32 @tgt() {
+entry:
+  %vs = call i32 @llvm.vscale.i32()
+  %len = mul i32 %vs, 2
+  %init_vec = insertelement <vscale x 2 x i32> poison, i32 0, i32 0
+  br label %for.cond
+
+for.cond:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %vec = phi <vscale x 2 x i32> [ %init_vec, %entry ], [ %new_vec, %for.body ]
+  %cmp = icmp slt i32 %i, %len
+  br i1 %cmp, label %for.body, label %exit
+
+for.body:
+  %new_vec = insertelement <vscale x 2 x i32> %vec, i32 %i, i32 %i
+  %inc = add i32 %i, 1
+  br label %for.cond
+
+exit:
+  %last_idx = sub i32 %len, 1
+  %result = extractelement <vscale x 2 x i32> %vec, i32 %last_idx
+  ret i32 %result
+}
+
+; TEST-ARGS: --vscale=3 -tgt-unroll=6

tests/alive-tv/vector/scalable/shufflevector-poison1.srctgt.ll

@@ -0,0 +1,11 @@
+define <vscale x 2 x i32> @src(<vscale x 2 x i32> %vec) {
+  %insert = insertelement <vscale x 2 x i32> %vec, i32 0, i32 0
+  %shuf = shufflevector <vscale x 2 x i32> %insert, <vscale x 2 x i32> poison, <vscale x 2 x i32> poison
+  ret <vscale x 2 x i32> %shuf
+}
+
+define <vscale x 2 x i32> @tgt(<vscale x 2 x i32> %vec) {
+  ret <vscale x 2 x i32> poison
+}
+
+; TEST-ARGS: --vscale=1

tests/alive-tv/vector/scalable/shufflevector-poison2.srctgt.ll

@@ -0,0 +1,11 @@
+define <vscale x 2 x i32> @src(<vscale x 2 x i32> %vec) {
+  %insert = insertelement <vscale x 2 x i32> %vec, i32 0, i32 0
+  %shuf = shufflevector <vscale x 2 x i32> %insert, <vscale x 2 x i32> poison, <vscale x 2 x i32> poison
+  ret <vscale x 2 x i32> %shuf
+}
+
+define <vscale x 2 x i32> @tgt(<vscale x 2 x i32> %vec) {
+  ret <vscale x 2 x i32> poison
+}
+
+; TEST-ARGS: --vscale=2