Fold 64-bit int operations

source/opt/const_folding_rules.cpp

@@ -21,6 +21,59 @@ namespace opt {
 namespace {
 constexpr uint32_t kExtractCompositeIdInIdx = 0;
 
+// Returns the value obtained by setting clearing the `number_of_bits` most
+// significant bits of `value`.
+uint64_t SignExtendValue(uint64_t value, uint32_t number_of_bits) {
+ if (number_of_bits == 64) return value;
+
+ uint64_t mask_for_sign_bit = 1ull << (number_of_bits - 1);
+ uint64_t mask_for_significant_bits = (mask_for_sign_bit << 1) - 1ull;
+ if (value & mask_for_sign_bit) {
+ // Set upper bits to 1
+ value |= ~mask_for_significant_bits;
+ } else {
+ // Clear the upper bits
+ value &= mask_for_significant_bits;
+ }
+ return value;
+}
+
+// Returns the value obtained from clearing the `number_of_bits` most
+// significant bits of `value`.
+uint64_t ClearUpperBits(uint64_t value, uint32_t number_of_bits) {
+ if (number_of_bits == 0) return value;
+
+ uint64_t mask_for_first_bit_to_clear = 1ull << (64 - number_of_bits);
+ uint64_t mask_for_bits_to_keep = mask_for_first_bit_to_clear - 1;
+ value &= mask_for_bits_to_keep;
+ return value;
+}
+
+// Returns a constant whose value is `value` and type is `type`. This constant
+// will be generated by `const_mgr`. The type must be a scalar integer type.
+const analysis::Constant* GenerateIntegerConstant(
+ const analysis::Integer* integer_type, uint64_t result,
+ analysis::ConstantManager* const_mgr) {
+ assert(integer_type != nullptr);
+
+ std::vector<uint32_t> words;
+ if (integer_type->width() == 64) {
+ // In the 64-bit case, two words are needed to represent the value.
+ words = {static_cast<uint32_t>(result),
+ static_cast<uint32_t>(result >> 32)};
+ } else {
+ // In all other cases, only a single word is needed.
+ assert(integer_type->width() <= 32);
+ if (integer_type->IsSigned()) {
+ result = SignExtendValue(result, integer_type->width());
+ } else {
+ result = ClearUpperBits(result, 64 - integer_type->width());
+ }
+ words = {static_cast<uint32_t>(result)};
+ }
+ return const_mgr->GetConstant(integer_type, words);
+}
+
 // Returns a constants with the value NaN of the given type. Only works for
 // 32-bit and 64-bit float point types. Returns |nullptr| if an error occurs.
 const analysis::Constant* GetNan(const analysis::Type* type,
@@ -676,7 +729,6 @@ ConstantFoldingRule FoldUnaryOp(UnaryScalarFoldingRule scalar_rule) {
  return [scalar_rule](IRContext* context, Instruction* inst,
  const std::vector<const analysis::Constant*>& constants)
  -> const analysis::Constant* {
-
  analysis::ConstantManager* const_mgr = context->get_constant_mgr();
  analysis::TypeManager* type_mgr = context->get_type_mgr();
  const analysis::Type* result_type = type_mgr->GetType(inst->type_id());
@@ -716,6 +768,63 @@ ConstantFoldingRule FoldUnaryOp(UnaryScalarFoldingRule scalar_rule) {
  };
 }
 
+// Returns a |ConstantFoldingRule| that folds binary scalar ops
+// using |scalar_rule| and unary vectors ops by applying
+// |scalar_rule| to the elements of the vector. The |ConstantFoldingRule|
+// that is returned assumes that |constants| contains 2 entries. If they are
+// not |nullptr|, then their type is either |Float| or |Integer| or a |Vector|
+// whose element type is |Float| or |Integer|.
+ConstantFoldingRule FoldBinaryOp(BinaryScalarFoldingRule scalar_rule) {
+ return [scalar_rule](IRContext* context, Instruction* inst,
+ const std::vector<const analysis::Constant*>& constants)
+ -> const analysis::Constant* {
+ analysis::ConstantManager* const_mgr = context->get_constant_mgr();
+ analysis::TypeManager* type_mgr = context->get_type_mgr();
+ const analysis::Type* result_type = type_mgr->GetType(inst->type_id());
+ const analysis::Vector* vector_type = result_type->AsVector();
+
+ const analysis::Constant* arg1 =
+ (inst->opcode() == spv::Op::OpExtInst) ? constants[1] : constants[0];
+ const analysis::Constant* arg2 =
+ (inst->opcode() == spv::Op::OpExtInst) ? constants[2] : constants[1];
+
+ if (arg1 == nullptr) {
+ return nullptr;
+ }
+ if (arg2 == nullptr) {
+ return nullptr;
+ }
+
+ if (vector_type != nullptr) {
+ std::vector<const analysis::Constant*> a_components;
+ std::vector<const analysis::Constant*> b_components;
+ std::vector<const analysis::Constant*> results_components;
+
+ a_components = arg1->GetVectorComponents(const_mgr);
+ b_components = arg2->GetVectorComponents(const_mgr);
+
+ // Fold each component of the vector.
+ for (uint32_t i = 0; i < a_components.size(); ++i) {
+ results_components.push_back(scalar_rule(vector_type->element_type(),
+ a_components[i],
+ b_components[i], const_mgr));
+ if (results_components[i] == nullptr) {
+ return nullptr;
+ }
+ }
+
+ // Build the constant object and return it.
+ std::vector<uint32_t> ids;
+ for (const analysis::Constant* member : results_components) {
+ ids.push_back(const_mgr->GetDefiningInstruction(member)->result_id());
+ }
+ return const_mgr->GetConstant(vector_type, ids);
+ } else {
+ return scalar_rule(result_type, arg1, arg2, const_mgr);
+ }
+ };
+}
+
 // Returns a |ConstantFoldingRule| that folds unary floating point scalar ops
 // using |scalar_rule| and unary float point vectors ops by applying
 // |scalar_rule| to the elements of the vector. The |ConstantFoldingRule|
@@ -1587,6 +1696,68 @@ BinaryScalarFoldingRule FoldFTranscendentalBinary(double (*fp)(double,
  return nullptr;
  };
 }
+
+enum Sign { Signed, Unsigned };
+
+// Returns a BinaryScalarFoldingRule that applies `op` to the scalars.
+// The `signedness` is used to determine if the operands should be interpreted
+// as signed or unsigned. If the operands are signed, the will be sign extended
+// before the value is passed to `op`. Otherwise the values will be zero
+// extended.
+template <Sign signedness>
+BinaryScalarFoldingRule FoldBinaryIntegerOperation(uint64_t (*op)(uint64_t,
+ uint64_t)) {
+ return
+ [op](const analysis::Type* result_type, const analysis::Constant* a,
+ const analysis::Constant* b,
+ analysis::ConstantManager* const_mgr) -> const analysis::Constant* {
+ assert(result_type != nullptr && a != nullptr);
+ const analysis::Integer* integer_type = a->type()->AsInteger();
+ assert(integer_type != nullptr);
+ assert(integer_type == result_type->AsInteger());
+ assert(integer_type == b->type()->AsInteger());
+
+ // In SPIR-V, the signedness of the operands is determined by the
+ // opcode, and not by the type of the operands. This is why we use the
+ // template argument to determine how to interpret the operands.
+ uint64_t ia = (signedness == Signed ? a->GetSignExtendedValue()
+ : a->GetZeroExtendedValue());
+ uint64_t ib = (signedness == Signed ? b->GetSignExtendedValue()
+ : b->GetZeroExtendedValue());
+ uint64_t result = op(ia, ib);
+
+ const analysis::Constant* result_constant =
+ GenerateIntegerConstant(integer_type, result, const_mgr);
+ return result_constant;
+ };
+}
+
+// A scalar folding rule that foles OpSConvert.
+const analysis::Constant* FoldScalarSConvert(
+ const analysis::Type* result_type, const analysis::Constant* a,
+ analysis::ConstantManager* const_mgr) {
+ assert(a);
+ const analysis::Integer* integer_type = result_type->AsInteger();
+ assert(integer_type && "The result type of an SConvert");
+ int64_t value = a->GetSignExtendedValue();
+ return GenerateIntegerConstant(integer_type, value, const_mgr);
+}
+
+// A scalar folding rule that foles OpSConvert.
+const analysis::Constant* FoldScalarUConvert(
+ const analysis::Type* result_type, const analysis::Constant* a,
+ analysis::ConstantManager* const_mgr) {
+ assert(a);
+ const analysis::Integer* integer_type = result_type->AsInteger();
+ assert(integer_type && "The result type of an SConvert");
+ uint64_t value = a->GetZeroExtendedValue();
+
+ // If the operand was an unsigned value with less than 32-bit, it would have
+ // been sign extended earlier, and we need to clear those bits.
+ auto* operand_type = a->type()->AsInteger();
+ value = ClearUpperBits(value, 64 - operand_type->width());
+ return GenerateIntegerConstant(integer_type, value, const_mgr);
+}
 } // namespace
 
 void ConstantFoldingRules::AddFoldingRules() {
@@ -1604,6 +1775,8 @@ void ConstantFoldingRules::AddFoldingRules() {
  rules_[spv::Op::OpConvertFToU].push_back(FoldFToI());
  rules_[spv::Op::OpConvertSToF].push_back(FoldIToF());
  rules_[spv::Op::OpConvertUToF].push_back(FoldIToF());
+ rules_[spv::Op::OpSConvert].push_back(FoldUnaryOp(FoldScalarSConvert));
+ rules_[spv::Op::OpUConvert].push_back(FoldUnaryOp(FoldScalarUConvert));
 
  rules_[spv::Op::OpDot].push_back(FoldOpDotWithConstants());
  rules_[spv::Op::OpFAdd].push_back(FoldFAdd());
@@ -1662,6 +1835,46 @@ void ConstantFoldingRules::AddFoldingRules() {
  rules_[spv::Op::OpSNegate].push_back(FoldSNegate());
  rules_[spv::Op::OpQuantizeToF16].push_back(FoldQuantizeToF16());
 
+ rules_[spv::Op::OpIAdd].push_back(
+ FoldBinaryOp(FoldBinaryIntegerOperation<Unsigned>(
+ [](uint64_t a, uint64_t b) { return a + b; })));
+ rules_[spv::Op::OpISub].push_back(
+ FoldBinaryOp(FoldBinaryIntegerOperation<Unsigned>(
+ [](uint64_t a, uint64_t b) { return a - b; })));
+ rules_[spv::Op::OpIMul].push_back(
+ FoldBinaryOp(FoldBinaryIntegerOperation<Unsigned>(
+ [](uint64_t a, uint64_t b) { return a * b; })));
+ rules_[spv::Op::OpUDiv].push_back(
+ FoldBinaryOp(FoldBinaryIntegerOperation<Unsigned>(
+ [](uint64_t a, uint64_t b) { return (b != 0 ? a / b : 0); })));
+ rules_[spv::Op::OpSDiv].push_back(FoldBinaryOp(
+ FoldBinaryIntegerOperation<Signed>([](uint64_t a, uint64_t b) {
+ return (b != 0 ? static_cast<uint64_t>(static_cast<int64_t>(a) /
+ static_cast<int64_t>(b))
+ : 0);
+ })));
+ rules_[spv::Op::OpUMod].push_back(
+ FoldBinaryOp(FoldBinaryIntegerOperation<Unsigned>(
+ [](uint64_t a, uint64_t b) { return (b != 0 ? a % b : 0); })));
+
+ rules_[spv::Op::OpSRem].push_back(FoldBinaryOp(
+ FoldBinaryIntegerOperation<Signed>([](uint64_t a, uint64_t b) {
+ return (b != 0 ? static_cast<uint64_t>(static_cast<int64_t>(a) %
+ static_cast<int64_t>(b))
+ : 0);
+ })));
+
+ rules_[spv::Op::OpSMod].push_back(FoldBinaryOp(
+ FoldBinaryIntegerOperation<Signed>([](uint64_t a, uint64_t b) {
+ if (b == 0) return static_cast<uint64_t>(0ull);
+
+ int64_t signed_a = static_cast<int64_t>(a);
+ int64_t signed_b = static_cast<int64_t>(b);
+ int64_t result = signed_a % signed_b;
+ if ((signed_b < 0) != (result < 0)) result += signed_b;
+ return static_cast<uint64_t>(result);
+ })));
+
  // Add rules for GLSLstd450
  FeatureManager* feature_manager = context_->get_feature_mgr();
  uint32_t ext_inst_glslstd450_id =