/*========================== begin_copyright_notice ============================

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#include "Compiler/IGCPassSupport.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Constants.h"
#include "llvm/Pass.h"
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;

namespace {

    class PruneUnusedArguments : public ModulePass {
    public:
        static char ID;
        PruneUnusedArguments();
        bool runOnModule(Module& M) override;
        void getAnalysisUsage(AnalysisUsage& AU) const override {
            AU.addRequired<CallGraphWrapperPass>();
            AU.addPreserved<CallGraphWrapperPass>();
            AU.setPreservesCFG();
        }
    };

} // namespace

namespace IGC {
    llvm::ModulePass* createPruneUnusedArgumentsPass() {
        return new PruneUnusedArguments();
    }
} // namespace IGC

IGC_INITIALIZE_PASS_BEGIN(PruneUnusedArguments, "PruneUnusedArguments", "PruneUnusedArguments", false, false)
IGC_INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
IGC_INITIALIZE_PASS_END(PruneUnusedArguments, "PruneUnusedArguments", "PruneUnusedArguments", false, false)

char PruneUnusedArguments::ID = 0;

/// Optimize unused arguments, all indirectly unused arguments will be replaced
/// with undef and vISA emission will skip these copies. This is in fact a
/// cleanup to the unification passes, which run early in the pipeline. With
/// subroutines, this issue becomes worse.
///
/// Note that (1) we do not delete those dead arguments since there may have
/// metadata on those functions; (2) the way we do private memory resolution may
/// introduce uses to r0 or payloadHeader(!!). For this reason, we run this pass
/// after private meomry resolution and no other pass should introduce such uses
/// after. If r0 is given by an intrinsic, then this is not an issue by design.
///
bool PruneUnusedArguments::runOnModule(Module& M) {
    CallGraph& CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();

    // Visit functions in a post order DFS, i.e. reversed topological ordering. If
    // an argument is not used then let its callers pass undef.
    bool Changed = false;
    for (auto I = po_begin(CG.getExternalCallingNode()),
        E = po_end(CG.getExternalCallingNode());
        I != E; ++I) {
        auto CGNode = *I;
        // Skip external and indirect nodes.
        if (auto F = CGNode->getFunction()) {
            // Conservatively declarations use all arguments.
            if (F->isDeclaration())
                continue;
            // Ignore externally linked functions
            if (F->hasFnAttribute("referenced-indirectly") ||
                F->hasFnAttribute("invoke_simd_target"))
                continue;

            // Collect unused arguments and their indices.
            SmallVector<std::pair<Argument*, unsigned>, 8> UnusedArgs;
            unsigned Index = 0;
            for (auto& Arg : F->args()) {
                if (Arg.use_empty())
                    UnusedArgs.push_back(std::make_pair(&Arg, Index));
                Index++;
            }

            if (UnusedArgs.empty())
                continue;

            // Update call sites.
            for (auto U : F->users()) {
                CallInst* CI = dyn_cast<CallInst>(U);
                if (!CI)
                    continue;
                for (const auto& Item : UnusedArgs) {
                    auto Arg = Item.first;
                    auto Index = Item.second;
                    if (!isa<UndefValue>(CI->getArgOperand(Index))) {
                        CI->setArgOperand(Index, UndefValue::get(Arg->getType()));
                        Changed = true;
                    }
                }
            }
        }
    }

    return Changed;
}

PruneUnusedArguments::PruneUnusedArguments() : ModulePass(ID) {
    initializePruneUnusedArgumentsPass(*PassRegistry::getPassRegistry());
}
