2021-08-14 00:04:46 +03:00
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
|
|
// //
|
|
|
|
// DxilPIXAddTidToAmplificationShaderPayload.cpp //
|
|
|
|
// Copyright (C) Microsoft Corporation. All rights reserved. //
|
|
|
|
// This file is distributed under the University of Illinois Open Source //
|
|
|
|
// License. See LICENSE.TXT for details. //
|
|
|
|
// //
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
|
|
|
|
|
|
#include "dxc/DXIL/DxilOperations.h"
|
|
|
|
#include "dxc/DXIL/DxilUtil.h"
|
|
|
|
|
|
|
|
#include "dxc/DXIL/DxilInstructions.h"
|
|
|
|
#include "dxc/DXIL/DxilModule.h"
|
|
|
|
#include "dxc/DxilPIXPasses/DxilPIXPasses.h"
|
|
|
|
|
|
|
|
#include "llvm/IR/InstIterator.h"
|
|
|
|
#include "llvm/IR/PassManager.h"
|
|
|
|
#include "llvm/Support/FormattedStream.h"
|
|
|
|
#include "llvm/Transforms/Utils/Local.h"
|
|
|
|
|
|
|
|
#include "PixPassHelpers.h"
|
|
|
|
|
|
|
|
using namespace llvm;
|
|
|
|
using namespace hlsl;
|
|
|
|
using namespace PIXPassHelpers;
|
|
|
|
|
|
|
|
class DxilPIXAddTidToAmplificationShaderPayload : public ModulePass {
|
|
|
|
uint32_t m_DispatchArgumentY = 1;
|
|
|
|
uint32_t m_DispatchArgumentZ = 1;
|
|
|
|
|
|
|
|
public:
|
|
|
|
static char ID; // Pass identification, replacement for typeid
|
|
|
|
DxilPIXAddTidToAmplificationShaderPayload() : ModulePass(ID) {}
|
2023-09-19 15:49:22 +03:00
|
|
|
StringRef getPassName() const override {
|
|
|
|
return "DXIL Add flat thread id to payload from AS to MS";
|
|
|
|
}
|
2021-08-14 00:04:46 +03:00
|
|
|
bool runOnModule(Module &M) override;
|
|
|
|
void applyOptions(PassOptions O) override;
|
|
|
|
};
|
|
|
|
|
|
|
|
void DxilPIXAddTidToAmplificationShaderPayload::applyOptions(PassOptions O) {
|
|
|
|
GetPassOptionUInt32(O, "dispatchArgY", &m_DispatchArgumentY, 1);
|
|
|
|
GetPassOptionUInt32(O, "dispatchArgZ", &m_DispatchArgumentZ, 1);
|
|
|
|
}
|
|
|
|
|
2023-09-19 15:49:22 +03:00
|
|
|
void AddValueToExpandedPayload(OP *HlslOP, llvm::IRBuilder<> &B,
|
|
|
|
ExpandedStruct &expanded,
|
|
|
|
AllocaInst *NewStructAlloca,
|
|
|
|
unsigned int expandedValueIndex, Value *value) {
|
2021-09-22 05:29:07 +03:00
|
|
|
Constant *Zero32Arg = HlslOP->GetU32Const(0);
|
|
|
|
SmallVector<Value *, 2> IndexToAppendedValue;
|
|
|
|
IndexToAppendedValue.push_back(Zero32Arg);
|
|
|
|
IndexToAppendedValue.push_back(HlslOP->GetU32Const(expandedValueIndex));
|
2023-09-19 15:49:22 +03:00
|
|
|
auto *PointerToEmbeddedNewValue = B.CreateInBoundsGEP(
|
|
|
|
expanded.ExpandedPayloadStructType, NewStructAlloca, IndexToAppendedValue,
|
|
|
|
"PointerToEmbeddedNewValue" + std::to_string(expandedValueIndex));
|
2021-09-22 05:29:07 +03:00
|
|
|
B.CreateStore(value, PointerToEmbeddedNewValue);
|
|
|
|
}
|
|
|
|
|
2021-08-14 00:04:46 +03:00
|
|
|
bool DxilPIXAddTidToAmplificationShaderPayload::runOnModule(Module &M) {
|
|
|
|
|
|
|
|
DxilModule &DM = M.GetOrCreateDxilModule();
|
|
|
|
LLVMContext &Ctx = M.getContext();
|
|
|
|
OP *HlslOP = DM.GetOP();
|
|
|
|
|
2023-09-19 15:49:22 +03:00
|
|
|
Type *OriginalPayloadStructPointerType = nullptr;
|
|
|
|
Type *OriginalPayloadStructType = nullptr;
|
2021-08-14 00:04:46 +03:00
|
|
|
ExpandedStruct expanded;
|
2023-09-19 15:49:22 +03:00
|
|
|
llvm::Function *entryFunction = PIXPassHelpers::GetEntryFunction(DM);
|
|
|
|
for (inst_iterator I = inst_begin(entryFunction), E = inst_end(entryFunction);
|
2021-08-14 00:04:46 +03:00
|
|
|
I != E; ++I) {
|
2023-09-19 15:49:22 +03:00
|
|
|
if (auto *Instr = llvm::cast<Instruction>(&*I)) {
|
|
|
|
if (hlsl::OP::IsDxilOpFuncCallInst(Instr,
|
|
|
|
hlsl::OP::OpCode::DispatchMesh)) {
|
|
|
|
DxilInst_DispatchMesh DispatchMesh(Instr);
|
|
|
|
OriginalPayloadStructPointerType =
|
|
|
|
DispatchMesh.get_payload()->getType();
|
|
|
|
OriginalPayloadStructType =
|
|
|
|
OriginalPayloadStructPointerType->getPointerElementType();
|
|
|
|
expanded = ExpandStructType(Ctx, OriginalPayloadStructType);
|
2021-08-14 00:04:46 +03:00
|
|
|
}
|
2023-09-19 15:49:22 +03:00
|
|
|
}
|
2021-08-14 00:04:46 +03:00
|
|
|
}
|
|
|
|
|
2023-09-19 15:49:22 +03:00
|
|
|
AllocaInst *OldStructAlloca = nullptr;
|
|
|
|
AllocaInst *NewStructAlloca = nullptr;
|
|
|
|
std::vector<AllocaInst *> allocasOfPayloadType;
|
|
|
|
for (inst_iterator I = inst_begin(entryFunction), E = inst_end(entryFunction);
|
|
|
|
I != E; ++I) {
|
|
|
|
auto *Inst = &*I;
|
|
|
|
if (llvm::isa<AllocaInst>(Inst)) {
|
|
|
|
auto *Alloca = llvm::cast<AllocaInst>(Inst);
|
|
|
|
if (Alloca->getType() == OriginalPayloadStructPointerType) {
|
|
|
|
allocasOfPayloadType.push_back(Alloca);
|
2021-08-14 00:04:46 +03:00
|
|
|
}
|
2023-09-19 15:49:22 +03:00
|
|
|
}
|
2021-08-14 00:04:46 +03:00
|
|
|
}
|
2023-09-19 15:49:22 +03:00
|
|
|
for (auto &Alloca : allocasOfPayloadType) {
|
2021-08-14 00:04:46 +03:00
|
|
|
OldStructAlloca = Alloca;
|
|
|
|
llvm::IRBuilder<> B(Alloca->getContext());
|
|
|
|
NewStructAlloca = B.CreateAlloca(expanded.ExpandedPayloadStructType,
|
|
|
|
HlslOP->GetU32Const(1), "NewPayload");
|
|
|
|
NewStructAlloca->setAlignment(Alloca->getAlignment());
|
|
|
|
NewStructAlloca->insertAfter(Alloca);
|
|
|
|
|
|
|
|
ReplaceAllUsesOfInstructionWithNewValueAndDeleteInstruction(
|
|
|
|
Alloca, NewStructAlloca, expanded.ExpandedPayloadStructType);
|
|
|
|
}
|
|
|
|
|
2023-09-19 15:49:22 +03:00
|
|
|
auto F = HlslOP->GetOpFunc(DXIL::OpCode::DispatchMesh,
|
|
|
|
expanded.ExpandedPayloadStructPtrType);
|
2021-08-14 00:04:46 +03:00
|
|
|
for (auto FI = F->user_begin(); FI != F->user_end();) {
|
2023-09-19 15:49:22 +03:00
|
|
|
auto *FunctionUser = *FI++;
|
|
|
|
auto *UserInstruction = llvm::cast<Instruction>(FunctionUser);
|
|
|
|
DxilInst_DispatchMesh DispatchMesh(UserInstruction);
|
|
|
|
|
|
|
|
llvm::IRBuilder<> B(UserInstruction);
|
|
|
|
|
|
|
|
auto ThreadIdFunc =
|
|
|
|
HlslOP->GetOpFunc(DXIL::OpCode::ThreadId, Type::getInt32Ty(Ctx));
|
|
|
|
Constant *Opcode = HlslOP->GetU32Const((unsigned)DXIL::OpCode::ThreadId);
|
|
|
|
Constant *Zero32Arg = HlslOP->GetU32Const(0);
|
|
|
|
Constant *One32Arg = HlslOP->GetU32Const(1);
|
|
|
|
Constant *Two32Arg = HlslOP->GetU32Const(2);
|
|
|
|
|
|
|
|
auto ThreadIdX =
|
|
|
|
B.CreateCall(ThreadIdFunc, {Opcode, Zero32Arg}, "ThreadIdX");
|
|
|
|
auto ThreadIdY =
|
|
|
|
B.CreateCall(ThreadIdFunc, {Opcode, One32Arg}, "ThreadIdY");
|
|
|
|
auto ThreadIdZ =
|
|
|
|
B.CreateCall(ThreadIdFunc, {Opcode, Two32Arg}, "ThreadIdZ");
|
|
|
|
|
|
|
|
auto *XxY =
|
|
|
|
B.CreateMul(ThreadIdX, HlslOP->GetU32Const(m_DispatchArgumentY));
|
|
|
|
auto *XplusY = B.CreateAdd(ThreadIdY, XxY);
|
|
|
|
auto *XYxZ = B.CreateMul(XplusY, HlslOP->GetU32Const(m_DispatchArgumentZ));
|
|
|
|
auto *XYZ = B.CreateAdd(ThreadIdZ, XYxZ);
|
|
|
|
|
|
|
|
AddValueToExpandedPayload(HlslOP, B, expanded, NewStructAlloca,
|
|
|
|
OriginalPayloadStructType->getStructNumElements(),
|
|
|
|
XYZ);
|
|
|
|
AddValueToExpandedPayload(
|
|
|
|
HlslOP, B, expanded, NewStructAlloca,
|
|
|
|
OriginalPayloadStructType->getStructNumElements() + 1,
|
|
|
|
DispatchMesh.get_threadGroupCountY());
|
|
|
|
AddValueToExpandedPayload(
|
|
|
|
HlslOP, B, expanded, NewStructAlloca,
|
|
|
|
OriginalPayloadStructType->getStructNumElements() + 2,
|
|
|
|
DispatchMesh.get_threadGroupCountZ());
|
2021-08-14 00:04:46 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
DM.ReEmitDxilResources();
|
|
|
|
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
|
|
|
char DxilPIXAddTidToAmplificationShaderPayload::ID = 0;
|
|
|
|
|
2023-09-19 15:49:22 +03:00
|
|
|
ModulePass *llvm::createDxilPIXAddTidToAmplificationShaderPayloadPass() {
|
|
|
|
return new DxilPIXAddTidToAmplificationShaderPayload();
|
|
|
|
}
|
2021-08-14 00:04:46 +03:00
|
|
|
|
2023-09-19 15:49:22 +03:00
|
|
|
INITIALIZE_PASS(DxilPIXAddTidToAmplificationShaderPayload,
|
|
|
|
"hlsl-dxil-PIX-add-tid-to-as-payload",
|
|
|
|
"HLSL DXIL Add flat thread id to payload from AS to MS", false,
|
|
|
|
false)
|