[HLSL] Raise Diag for Invalid CounterDirection

Author: V-FEXrt

llvm/include/llvm/Analysis/DXILResource.h

@@ -451,8 +451,10 @@ ModulePass *createDXILResourceTypeWrapperPassPass();
 //===----------------------------------------------------------------------===//
 
 class DXILResourceMap {
+  using CallMapTy = DenseMap<CallInst *, unsigned>;
+
   SmallVector<dxil::ResourceInfo> Infos;
-  DenseMap<CallInst *, unsigned> CallMap;
+  CallMapTy CallMap;
   unsigned FirstUAV = 0;
   unsigned FirstCBuffer = 0;
   unsigned FirstSampler = 0;
@@ -532,6 +534,21 @@ class DXILResourceMap {
     return make_range(sampler_begin(), sampler_end());
   }
 
+  struct call_iterator
+      : iterator_adaptor_base<call_iterator, CallMapTy::iterator> {
+    call_iterator() = default;
+    call_iterator(CallMapTy::iterator Iter)
+        : call_iterator::iterator_adaptor_base(std::move(Iter)) {}
+
+    CallInst *operator*() const { return I->first; }
+  };
+
+  call_iterator call_begin() { return call_iterator(CallMap.begin()); }
+  call_iterator call_end() { return call_iterator(CallMap.end()); }
+  iterator_range<call_iterator> calls() {
+    return make_range(call_begin(), call_end());
+  }
+
   void print(raw_ostream &OS, DXILResourceTypeMap &DRTM,
              const DataLayout &DL) const;
 

llvm/lib/Target/DirectX/CMakeLists.txt

@@ -28,6 +28,7 @@ add_llvm_target(DirectXCodeGen
   DXILIntrinsicExpansion.cpp
   DXILOpBuilder.cpp
   DXILOpLowering.cpp
+  DXILPostOptimizationValidation.cpp
   DXILPrepare.cpp
   DXILPrettyPrinter.cpp
   DXILResourceAccess.cpp

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -0,0 +1,119 @@
+//===- DXILPostOptimizationValidation.cpp - Opt DXIL validation ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "DXILPostOptimizationValidation.h"
+#include "DXILConstants.h"
+#include "DXILIntrinsicExpansion.h"
+#include "DXILOpBuilder.h"
+#include "DXILShaderFlags.h"
+#include "DirectX.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/DXILMetadataAnalysis.h"
+#include "llvm/Analysis/DXILResource.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/ErrorHandling.h"
+#include <cstdio>
+
+#define DEBUG_TYPE "dxil-post-optimization-validation"
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+namespace {
+class DXILValidator {
+  Module &M;
+  DXILResourceMap &DRM;
+
+public:
+  DXILValidator(Module &M, DXILResourceMap &DRM) : M(M), DRM(DRM) {}
+
+  void validate() {
+    for (const auto &UAV : DRM.uavs()) {
+      if (UAV.CounterDirection != ResourceCounterDirection::Invalid)
+        continue;
+
+      CallInst *ResourceHandle = nullptr;
+      for (CallInst *MaybeHandle : DRM.calls()) {
+        if (*DRM.find(MaybeHandle) == UAV) {
+          ResourceHandle = MaybeHandle;
+          break;
+        }
+      }
+
+      StringRef Message =
+          "RWStructuredBuffers may increment or decrement their "
+          "counters, but not both.";
+      for (const auto &U : ResourceHandle->users()) {
+        const CallInst *CI = dyn_cast<CallInst>(U);
+        if (!CI && CI->getIntrinsicID() != Intrinsic::dx_resource_updatecounter)
+          continue;
+
+        M.getContext().diagnose(DiagnosticInfoGenericWithLoc(
+            Message, *CI->getFunction(), CI->getDebugLoc()));
+      }
+    }
+  }
+};
+} // namespace
+
+PreservedAnalyses
+DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {
+  DXILResourceMap &DRM = MAM.getResult<DXILResourceAnalysis>(M);
+
+  DXILValidator(M, DRM).validate();
+  return PreservedAnalyses::all();
+}
+
+namespace {
+class DXILPostOptimizationValidationLegacy : public ModulePass {
+public:
+  bool runOnModule(Module &M) override {
+    DXILResourceMap &DRM =
+        getAnalysis<DXILResourceWrapperPass>().getResourceMap();
+
+    DXILValidator(M, DRM).validate();
+
+    return false;
+  }
+  StringRef getPassName() const override {
+    return "DXIL Post Optimization Validation";
+  }
+  DXILPostOptimizationValidationLegacy() : ModulePass(ID) {}
+
+  static char ID; // Pass identification.
+  void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
+    AU.addRequired<DXILResourceWrapperPass>();
+    AU.addPreserved<DXILResourceWrapperPass>();
+    AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
+    AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+  }
+};
+char DXILPostOptimizationValidationLegacy::ID = 0;
+} // end anonymous namespace
+
+INITIALIZE_PASS_BEGIN(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
+                      "DXIL Post Optimization Validation", false, false)
+INITIALIZE_PASS_DEPENDENCY(DXILResourceTypeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DXILResourceWrapperPass)
+INITIALIZE_PASS_END(DXILPostOptimizationValidationLegacy, DEBUG_TYPE,
+                    "DXIL Post Optimization Validation", false, false)
+
+ModulePass *llvm::createDXILPostOptimizationValidationLegacyPass() {
+  return new DXILPostOptimizationValidationLegacy();
+}

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.h

@@ -0,0 +1,28 @@
+//===- DXILPostOptimizationValidation.h - Opt DXIL Validations -*- C++ -*--===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// \file Pass for validating DXIL after lowering and optimizations are applied.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_DIRECTX_DXILPOSTOPTIMIZATIONVALIDATION_H
+#define LLVM_LIB_TARGET_DIRECTX_DXILPOSTOPTIMIZATIONVALIDATION_H
+
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+
+class DXILPostOptimizationValidation
+    : public PassInfoMixin<DXILPostOptimizationValidation> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_LIB_TARGET_DIRECTX_DXILPOSTOPTIMIZATIONVALIDATION_H

llvm/lib/Target/DirectX/DirectX.h

@@ -90,6 +90,12 @@ ModulePass *createDXILPrettyPrinterLegacyPass(raw_ostream &OS);
 /// Initializer for DXILPrettyPrinter.
 void initializeDXILPrettyPrinterLegacyPass(PassRegistry &);
 
+/// Initializer for DXILPostOptimizationValidation.
+void initializeDXILPostOptimizationValidationLegacyPass(PassRegistry &);
+
+/// Pass to lowering LLVM intrinsic call to DXIL op function call.
+ModulePass *createDXILPostOptimizationValidationLegacyPass();
+
 /// Initializer for dxil::ShaderFlagsAnalysisWrapper pass.
 void initializeShaderFlagsAnalysisWrapperPass(PassRegistry &);
 

llvm/lib/Target/DirectX/DirectXPassRegistry.def

@@ -30,6 +30,7 @@ MODULE_PASS("dxil-intrinsic-expansion", DXILIntrinsicExpansion())
 MODULE_PASS("dxil-op-lower", DXILOpLowering())
 MODULE_PASS("dxil-pretty-printer", DXILPrettyPrinterPass(dbgs()))
 MODULE_PASS("dxil-translate-metadata", DXILTranslateMetadata())
+MODULE_PASS("dxil-post-optimization-validation", DXILPostOptimizationValidation())
 // TODO: rename to print<foo> after NPM switch
 MODULE_PASS("print-dx-shader-flags", dxil::ShaderFlagsAnalysisPrinter(dbgs()))
 MODULE_PASS("print<dxil-root-signature>", dxil::RootSignatureAnalysisPrinter(dbgs()))

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

@@ -19,6 +19,7 @@
 #include "DXILIntrinsicExpansion.h"
 #include "DXILLegalizePass.h"
 #include "DXILOpLowering.h"
+#include "DXILPostOptimizationValidation.h"
 #include "DXILPrettyPrinter.h"
 #include "DXILResourceAccess.h"
 #include "DXILRootSignature.h"
@@ -63,6 +64,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {
   initializeDXILOpLoweringLegacyPass(*PR);
   initializeDXILResourceAccessLegacyPass(*PR);
   initializeDXILTranslateMetadataLegacyPass(*PR);
+  initializeDXILPostOptimizationValidationLegacyPass(*PR);
   initializeShaderFlagsAnalysisWrapperPass(*PR);
   initializeRootSignatureAnalysisWrapperPass(*PR);
   initializeDXILFinalizeLinkageLegacyPass(*PR);
@@ -110,6 +112,7 @@ class DirectXPassConfig : public TargetPassConfig {
     addPass(createDXILForwardHandleAccessesLegacyPass());
     addPass(createDXILLegalizeLegacyPass());
     addPass(createDXILTranslateMetadataLegacyPass());
+    addPass(createDXILPostOptimizationValidationLegacyPass());
     addPass(createDXILOpLoweringLegacyPass());
     addPass(createDXILPrepareModulePass());
   }

llvm/test/CodeGen/DirectX/llc-pipeline.ll

@@ -28,6 +28,7 @@
 ; CHECK-NEXT:   DXIL Module Metadata analysis
 ; CHECK-NEXT:   DXIL Shader Flag Analysis
 ; CHECK-NEXT:   DXIL Translate Metadata
+; CHECK-NEXT:   DXIL Post Optimization Validation
 ; CHECK-NEXT:   DXIL Op Lowering
 ; CHECK-NEXT:   DXIL Prepare Module
 

llvm/test/CodeGen/DirectX/resource_counter_error.ll

@@ -0,0 +1,10 @@
+; RUN: not opt -S -passes='dxil-post-optimization-validation' -mtriple=dxil-pc-shadermodel6.3-library %s 2>&1 | FileCheck %s
+; CHECK: RWStructuredBuffers may increment or decrement their counters, but not both.
+
+define void @inc_and_dec() {
+entry:
+  %handle = call target("dx.RawBuffer", float, 1, 0) @llvm.dx.resource.handlefrombinding(i32 1, i32 2, i32 3, i32 4, i1 false)
+  call i32 @llvm.dx.resource.updatecounter(target("dx.RawBuffer", float, 1, 0) %handle, i8 -1)
+  call i32 @llvm.dx.resource.updatecounter(target("dx.RawBuffer", float, 1, 0) %handle, i8 1)
+  ret void
+}