TheRock/rocm-libraries/shared/rocroller/lib/source/ExpressionTransformations/LowerPRNG.cpp

// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include <rocRoller/Expression.hpp>

namespace rocRoller
{
    namespace Expression
    {

        struct LowerPRNGExpressionVisitor
        {
            template <CUnary Expr>
            ExpressionPtr operator()(Expr const& expr) const
            {
                Expr cpy = expr;
                cpy.arg  = call(expr.arg);
                return std::make_shared<Expression>(cpy);
            }

            template <CBinary Expr>
            ExpressionPtr operator()(Expr const& expr) const
            {
                Expr cpy = expr;
                cpy.lhs  = call(expr.lhs);
                cpy.rhs  = call(expr.rhs);
                return std::make_shared<Expression>(cpy);
            }

            template <CTernary Expr>
            ExpressionPtr operator()(Expr const& expr) const
            {
                Expr cpy = expr;
                cpy.lhs  = call(expr.lhs);
                cpy.r1hs = call(expr.r1hs);
                cpy.r2hs = call(expr.r2hs);
                return std::make_shared<Expression>(cpy);
            }

            template <CNary Expr>
            ExpressionPtr operator()(Expr const& expr) const
            {
                auto cpy = expr;
                std::ranges::for_each(cpy.operands, [this](auto& op) { op = call(op); });
                return std::make_shared<Expression>(std::move(cpy));
            }

            ExpressionPtr operator()(ScaledMatrixMultiply const& expr) const
            {
                ScaledMatrixMultiply cpy = expr;
                cpy.matA                 = call(expr.matA);
                cpy.matB                 = call(expr.matB);
                cpy.matC                 = call(expr.matC);
                cpy.scaleA               = call(expr.scaleA);
                cpy.scaleB               = call(expr.scaleB);
                return std::make_shared<Expression>(cpy);
            }

            ExpressionPtr operator()(RandomNumber const& expr) const
            {
                auto arg = call(expr.arg);

                // PRNG algorithm: ((seed << 1) ^ (((seed >> 31) & 1) ? 0xc5 : 0x00))
                ExpressionPtr one = literal(1u);

                ExpressionPtr lhs = std::make_shared<Expression>(ShiftL({arg, one}));

                ExpressionPtr shiftR
                    = std::make_shared<Expression>(LogicalShiftR({arg, literal(31u)}));
                ExpressionPtr bitAnd    = std::make_shared<Expression>(BitwiseAnd({shiftR, one}));
                ExpressionPtr predicate = std::make_shared<Expression>(Equal({bitAnd, one}));

                // Note: here we compute the `xor 0xc5` even though the value might not be selected.
                //       The reason of doing this instead of
                //          rhs = condition(predicate, 0xc5, 0x00)
                //          result = xor(value, rhs)
                //       is that `v_cndmask_b32_e64 v5, 0, 197, s[12:13]`  is not able to assemble
                //       (Error: literal operands are not supported).
                ExpressionPtr xorValue
                    = std::make_shared<Expression>(BitwiseXor({literal(197u), lhs})); // xor 0xc5

                return std::make_shared<Expression>(Conditional({predicate, xorValue, lhs}));
            }

            template <CValue Value>
            ExpressionPtr operator()(Value const& expr) const
            {
                return std::make_shared<Expression>(expr);
            }

            ExpressionPtr call(ExpressionPtr expr) const
            {
                if(!expr)
                    return expr;

                return std::visit(*this, *expr);
            }
        };

        /**
         *  Replace RandomNumber expression with equivalent expressions
         */
        ExpressionPtr lowerPRNG(ExpressionPtr expr, ContextPtr cxt)
        {
            if(cxt->targetArchitecture().HasCapability(GPUCapability::HasPRNG))
                return expr;

            auto visitor = LowerPRNGExpressionVisitor();
            return visitor.call(expr);
        }

    }
}