src/tint/lang/core/ir/transform/binary_polyfill.cc - tint - Git at Google

 // Copyright 2023 The Tint Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #include "src/tint/lang/core/ir/transform/binary_polyfill.h"

 #include <utility>

 #include "src/tint/lang/core/ir/builder.h"
 #include "src/tint/lang/core/ir/module.h"
 #include "src/tint/lang/core/ir/validator.h"

 using namespace tint::core::fluent_types;     // NOLINT
 using namespace tint::core::number_suffixes;  // NOLINT

 namespace tint::core::ir::transform {

 namespace {

 /// PIMPL state for the transform.
 struct State {
     /// The polyfill config.
     const BinaryPolyfillConfig& config;

     /// The IR module.
     Module& ir;

     /// The IR builder.
     Builder b{ir};

     /// The type manager.
     core::type::Manager& ty{ir.Types()};

     /// The symbol table.
     SymbolTable& sym{ir.symbols};

     /// Map from integer type to its divide helper function.
     Hashmap<const type::Type*, Function*, 4> int_div_helpers{};

     /// Map from integer type to its modulo helper function.
     Hashmap<const type::Type*, Function*, 4> int_mod_helpers{};

     /// Process the module.
     void Process() {
         // Find the binary instructions that need to be polyfilled.
         Vector<ir::Binary*, 64> worklist;
         for (auto* inst : ir.instructions.Objects()) {
             if (!inst->Alive()) {
                 continue;
             }
             if (auto* binary = inst->As<ir::Binary>()) {
                 switch (binary->Kind()) {
                     case ir::Binary::Kind::kDivide:
                     case ir::Binary::Kind::kModulo:
                         if (config.int_div_mod &&
                             binary->Result()->Type()->is_integer_scalar_or_vector()) {
                             worklist.Push(binary);
                         }
                         break;
                     case ir::Binary::Kind::kShiftLeft:
                     case ir::Binary::Kind::kShiftRight:
                         if (config.bitshift_modulo) {
                             worklist.Push(binary);
                         }
                         break;
                     default:
                         break;
                 }
             }
         }

         // Polyfill the binary instructions that we found.
         for (auto* binary : worklist) {
             ir::Value* replacement = nullptr;
             switch (binary->Kind()) {
                 case Binary::Kind::kDivide:
                 case Binary::Kind::kModulo:
                     replacement = IntDivMod(binary);
                     break;
                 case Binary::Kind::kShiftLeft:
                 case Binary::Kind::kShiftRight:
                     replacement = MaskShiftAmount(binary);
                     break;
                 default:
                     break;
             }
             TINT_ASSERT_OR_RETURN(replacement);

             if (replacement != binary->Result()) {
                 // Replace the old binary instruction result with the new value.
                 if (auto name = ir.NameOf(binary->Result())) {
                     ir.SetName(replacement, name);
                 }
                 binary->Result()->ReplaceAllUsesWith(replacement);
                 binary->Destroy();
             }
         }
     }

     /// Return a type with element type @p type that has the same number of vector components as
     /// @p match. If @p match is scalar just return @p type.
     /// @param el_ty the type to extend
     /// @param match the type to match the component count of
     /// @returns a type with the same number of vector components as @p match
     const core::type::Type* MatchWidth(const core::type::Type* el_ty,
                                        const core::type::Type* match) {
         if (auto* vec = match->As<core::type::Vector>()) {
             return ty.vec(el_ty, vec->Width());
         }
         return el_ty;
     }

     /// Return a constant that has the same number of vector components as @p match, each with the
     /// value @p element. If @p match is scalar just return @p element.
     /// @param element the value to extend
     /// @param match the type to match the component count of
     /// @returns a value with the same number of vector components as @p match
     ir::Constant* MatchWidth(ir::Constant* element, const core::type::Type* match) {
         if (auto* vec = match->As<core::type::Vector>()) {
             return b.Splat(MatchWidth(element->Type(), match), element, vec->Width());
         }
         return element;
     }

     /// Replace an integer divide or modulo with a call to helper function that prevents
     /// divide-by-zero and signed integer overflow.
     /// @param binary the binary instruction
     /// @returns the replacement value
     ir::Value* IntDivMod(ir::Binary* binary) {
         auto* result_ty = binary->Result()->Type();
         bool is_div = binary->Kind() == Binary::Kind::kDivide;
         bool is_signed = result_ty->is_signed_integer_scalar_or_vector();

         auto& helpers = is_div ? int_div_helpers : int_mod_helpers;
         auto* helper = helpers.GetOrCreate(result_ty, [&] {
             // Generate a name for the helper function.
             StringStream name;
             name << "tint_" << (is_div ? "div_" : "mod_");
             if (auto* vec = result_ty->As<type::Vector>()) {
                 name << "v" << vec->Width() << vec->type()->FriendlyName();
             } else {
                 name << result_ty->FriendlyName();
             }

             // Create the helper function.
             auto* func = b.Function(name.str(), result_ty);
             auto* lhs = b.FunctionParam("lhs", result_ty);
             auto* rhs = b.FunctionParam("rhs", result_ty);
             func->SetParams({lhs, rhs});
             b.Append(func->Block(), [&] {
                 // Generate constants for zero and one with types that match the result type.
                 ir::Constant* one = nullptr;
                 ir::Constant* zero = nullptr;
                 if (is_signed) {
                     one = MatchWidth(b.Constant(1_i), result_ty);
                     zero = MatchWidth(b.Constant(0_i), result_ty);
                 } else {
                     one = MatchWidth(b.Constant(1_u), result_ty);
                     zero = MatchWidth(b.Constant(0_u), result_ty);
                 }

                 // Select either the RHS or a constant one value if the RHS is zero.
                 // If this is a signed operation, we also check for `INT_MIN / -1`.
                 auto* bool_ty = MatchWidth(ty.bool_(), result_ty);
                 auto* cond = b.Equal(bool_ty, rhs, zero);
                 if (is_signed) {
                     auto* lowest = MatchWidth(b.Constant(i32::Lowest()), result_ty);
                     auto* minus_one = MatchWidth(b.Constant(-1_i), result_ty);
                     auto* lhs_is_lowest = b.Equal(bool_ty, lhs, lowest);
                     auto* rhs_is_minus_one = b.Equal(bool_ty, rhs, minus_one);
                     cond = b.Or(bool_ty, cond, b.And(bool_ty, lhs_is_lowest, rhs_is_minus_one));
                 }
                 auto* rhs_or_one = b.Call(result_ty, core::BuiltinFn::kSelect, rhs, one, cond);

                 if (binary->Kind() == Binary::Kind::kDivide) {
                     // Perform the divide with the modified RHS.
                     b.Return(func, b.Divide(result_ty, lhs, rhs_or_one)->Result());
                 } else if (binary->Kind() == Binary::Kind::kModulo) {
                     // Calculate the modulo manually, as modulo with negative operands is undefined
                     // behavior for many backends:
                     //   result = lhs - ((lhs / rhs_or_one) * rhs_or_one)
                     auto* whole = b.Divide(result_ty, lhs, rhs_or_one);
                     auto* remainder =
                         b.Subtract(result_ty, lhs, b.Multiply(result_ty, whole, rhs_or_one));
                     b.Return(func, remainder->Result());
                 }
             });
             return func;
         });

         /// Helper to splat a value to match the vector width of the result type if necessary.
         auto maybe_splat = [&](ir::Value* value) -> ir::Value* {
             if (value->Type()->Is<type::Scalar>() && result_ty->Is<core::type::Vector>()) {
                 return b.Construct(result_ty, value)->Result();
             }
             return value;
         };

         // Call the helper function, splatting the arguments to match the target vector width.
         Value* result = nullptr;
         b.InsertBefore(binary, [&] {
             auto* lhs = maybe_splat(binary->LHS());
             auto* rhs = maybe_splat(binary->RHS());
             result = b.Call(result_ty, helper, lhs, rhs)->Result();
         });
         return result;
     }

     /// Mask the RHS of a shift instruction to ensure it is modulo the bitwidth of the LHS.
     /// @param binary the binary instruction
     /// @returns the replacement value
     ir::Value* MaskShiftAmount(ir::Binary* binary) {
         auto* lhs = binary->LHS();
         auto* rhs = binary->RHS();
         auto* mask = b.Constant(u32(lhs->Type()->DeepestElement()->Size() * 8 - 1));
         auto* masked = b.And(rhs->Type(), rhs, MatchWidth(mask, rhs->Type()));
         masked->InsertBefore(binary);
         binary->SetOperand(ir::Binary::kRhsOperandOffset, masked->Result());
         return binary->Result();
     }
 };

 }  // namespace

 Result<SuccessType> BinaryPolyfill(Module& ir, const BinaryPolyfillConfig& config) {
     auto result = ValidateAndDumpIfNeeded(ir, "BinaryPolyfill transform");
     if (!result) {
         return result;
     }

     State{config, ir}.Process();

     return Success;
 }

 }  // namespace tint::core::ir::transform
	// Copyright 2023 The Tint Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#include "src/tint/lang/core/ir/transform/binary_polyfill.h"

	#include <utility>

	#include "src/tint/lang/core/ir/builder.h"
	#include "src/tint/lang/core/ir/module.h"
	#include "src/tint/lang/core/ir/validator.h"

	using namespace tint::core::fluent_types; // NOLINT
	using namespace tint::core::number_suffixes; // NOLINT

	namespace tint::core::ir::transform {

	namespace {

	/// PIMPL state for the transform.
	struct State {
	/// The polyfill config.
	const BinaryPolyfillConfig& config;

	/// The IR module.
	Module& ir;

	/// The IR builder.
	Builder b{ir};

	/// The type manager.
	core::type::Manager& ty{ir.Types()};

	/// The symbol table.
	SymbolTable& sym{ir.symbols};

	/// Map from integer type to its divide helper function.
	Hashmap<const type::Type, Function, 4> int_div_helpers{};

	/// Map from integer type to its modulo helper function.
	Hashmap<const type::Type, Function, 4> int_mod_helpers{};

	/// Process the module.
	void Process() {
	// Find the binary instructions that need to be polyfilled.
	Vector<ir::Binary*, 64> worklist;
	for (auto* inst : ir.instructions.Objects()) {
	if (!inst->Alive()) {
	continue;
	}
	if (auto* binary = inst->As<ir::Binary>()) {
	switch (binary->Kind()) {
	case ir::Binary::Kind::kDivide:
	case ir::Binary::Kind::kModulo:
	if (config.int_div_mod &&
	binary->Result()->Type()->is_integer_scalar_or_vector()) {
	worklist.Push(binary);
	}
	break;
	case ir::Binary::Kind::kShiftLeft:
	case ir::Binary::Kind::kShiftRight:
	if (config.bitshift_modulo) {
	worklist.Push(binary);
	}
	break;
	default:
	break;
	}
	}
	}

	// Polyfill the binary instructions that we found.
	for (auto* binary : worklist) {
	ir::Value* replacement = nullptr;
	switch (binary->Kind()) {
	case Binary::Kind::kDivide:
	case Binary::Kind::kModulo:
	replacement = IntDivMod(binary);
	break;
	case Binary::Kind::kShiftLeft:
	case Binary::Kind::kShiftRight:
	replacement = MaskShiftAmount(binary);
	break;
	default:
	break;
	}
	TINT_ASSERT_OR_RETURN(replacement);

	if (replacement != binary->Result()) {
	// Replace the old binary instruction result with the new value.
	if (auto name = ir.NameOf(binary->Result())) {
	ir.SetName(replacement, name);
	}
	binary->Result()->ReplaceAllUsesWith(replacement);
	binary->Destroy();
	}
	}
	}

	/// Return a type with element type @p type that has the same number of vector components as
	/// @p match. If @p match is scalar just return @p type.
	/// @param el_ty the type to extend
	/// @param match the type to match the component count of
	/// @returns a type with the same number of vector components as @p match
	const core::type::Type* MatchWidth(const core::type::Type* el_ty,
	const core::type::Type* match) {
	if (auto* vec = match->As<core::type::Vector>()) {
	return ty.vec(el_ty, vec->Width());
	}
	return el_ty;
	}

	/// Return a constant that has the same number of vector components as @p match, each with the
	/// value @p element. If @p match is scalar just return @p element.
	/// @param element the value to extend
	/// @param match the type to match the component count of
	/// @returns a value with the same number of vector components as @p match
	ir::Constant* MatchWidth(ir::Constant* element, const core::type::Type* match) {
	if (auto* vec = match->As<core::type::Vector>()) {
	return b.Splat(MatchWidth(element->Type(), match), element, vec->Width());
	}
	return element;
	}

	/// Replace an integer divide or modulo with a call to helper function that prevents
	/// divide-by-zero and signed integer overflow.
	/// @param binary the binary instruction
	/// @returns the replacement value
	ir::Value* IntDivMod(ir::Binary* binary) {
	auto* result_ty = binary->Result()->Type();
	bool is_div = binary->Kind() == Binary::Kind::kDivide;
	bool is_signed = result_ty->is_signed_integer_scalar_or_vector();

	auto& helpers = is_div ? int_div_helpers : int_mod_helpers;
	auto* helper = helpers.GetOrCreate(result_ty, [&] {
	// Generate a name for the helper function.
	StringStream name;
	name << "tint_" << (is_div ? "div_" : "mod_");
	if (auto* vec = result_ty->As<type::Vector>()) {
	name << "v" << vec->Width() << vec->type()->FriendlyName();
	} else {
	name << result_ty->FriendlyName();
	}

	// Create the helper function.
	auto* func = b.Function(name.str(), result_ty);
	auto* lhs = b.FunctionParam("lhs", result_ty);
	auto* rhs = b.FunctionParam("rhs", result_ty);
	func->SetParams({lhs, rhs});
	b.Append(func->Block(), [&] {
	// Generate constants for zero and one with types that match the result type.
	ir::Constant* one = nullptr;
	ir::Constant* zero = nullptr;
	if (is_signed) {
	one = MatchWidth(b.Constant(1_i), result_ty);
	zero = MatchWidth(b.Constant(0_i), result_ty);
	} else {
	one = MatchWidth(b.Constant(1_u), result_ty);
	zero = MatchWidth(b.Constant(0_u), result_ty);
	}

	// Select either the RHS or a constant one value if the RHS is zero.
	// If this is a signed operation, we also check for `INT_MIN / -1`.
	auto* bool_ty = MatchWidth(ty.bool_(), result_ty);
	auto* cond = b.Equal(bool_ty, rhs, zero);
	if (is_signed) {
	auto* lowest = MatchWidth(b.Constant(i32::Lowest()), result_ty);
	auto* minus_one = MatchWidth(b.Constant(-1_i), result_ty);
	auto* lhs_is_lowest = b.Equal(bool_ty, lhs, lowest);
	auto* rhs_is_minus_one = b.Equal(bool_ty, rhs, minus_one);
	cond = b.Or(bool_ty, cond, b.And(bool_ty, lhs_is_lowest, rhs_is_minus_one));
	}
	auto* rhs_or_one = b.Call(result_ty, core::BuiltinFn::kSelect, rhs, one, cond);

	if (binary->Kind() == Binary::Kind::kDivide) {
	// Perform the divide with the modified RHS.
	b.Return(func, b.Divide(result_ty, lhs, rhs_or_one)->Result());
	} else if (binary->Kind() == Binary::Kind::kModulo) {
	// Calculate the modulo manually, as modulo with negative operands is undefined
	// behavior for many backends:
	// result = lhs - ((lhs / rhs_or_one) * rhs_or_one)
	auto* whole = b.Divide(result_ty, lhs, rhs_or_one);
	auto* remainder =
	b.Subtract(result_ty, lhs, b.Multiply(result_ty, whole, rhs_or_one));
	b.Return(func, remainder->Result());
	}
	});
	return func;
	});

	/// Helper to splat a value to match the vector width of the result type if necessary.
	auto maybe_splat = [&](ir::Value* value) -> ir::Value* {
	if (value->Type()->Is<type::Scalar>() && result_ty->Is<core::type::Vector>()) {
	return b.Construct(result_ty, value)->Result();
	}
	return value;
	};

	// Call the helper function, splatting the arguments to match the target vector width.
	Value* result = nullptr;
	b.InsertBefore(binary, [&] {
	auto* lhs = maybe_splat(binary->LHS());
	auto* rhs = maybe_splat(binary->RHS());
	result = b.Call(result_ty, helper, lhs, rhs)->Result();
	});
	return result;
	}

	/// Mask the RHS of a shift instruction to ensure it is modulo the bitwidth of the LHS.
	/// @param binary the binary instruction
	/// @returns the replacement value
	ir::Value* MaskShiftAmount(ir::Binary* binary) {
	auto* lhs = binary->LHS();
	auto* rhs = binary->RHS();
	auto* mask = b.Constant(u32(lhs->Type()->DeepestElement()->Size() * 8 - 1));
	auto* masked = b.And(rhs->Type(), rhs, MatchWidth(mask, rhs->Type()));
	masked->InsertBefore(binary);
	binary->SetOperand(ir::Binary::kRhsOperandOffset, masked->Result());
	return binary->Result();
	}
	};

	} // namespace

	Result<SuccessType> BinaryPolyfill(Module& ir, const BinaryPolyfillConfig& config) {
	auto result = ValidateAndDumpIfNeeded(ir, "BinaryPolyfill transform");
	if (!result) {
	return result;
	}

	State{config, ir}.Process();

	return Success;
	}

	} // namespace tint::core::ir::transform