Fix dx raster codegen

include/luisa/backends/ext/dx_hdr_ext.hpp

@@ -0,0 +1,9 @@
+#pragma once
+#include <luisa/backends/ext/dx_hdr_ext_interface.h>
+#include <luisa/runtime/swapchain.h>
+#include <luisa/runtime/stream.h>
+namespace luisa::compute {
+Swapchain DXHDRExt::create_swapchain(const Stream &stream, const DXSwapchainOption &option) noexcept {
+    return Swapchain{stream.device(), create_swapchain(option, stream.handle())};
+}
+}// namespace luisa::compute

include/luisa/backends/ext/dx_hdr_ext_interface.h

@@ -0,0 +1,87 @@
+#pragma once
+#include <luisa/runtime/rhi/device_interface.h>
+
+namespace luisa::compute {
+class Swapchain;
+class Stream;
+class DXHDRExt : public DeviceExtension {
+public:
+    enum class DisplayCurve {
+        sRGB = 0,// The display expects an sRGB signal.
+        ST2084,  // The display expects an HDR10 signal.
+        None     // The display expects a linear signal.
+    };
+    enum class SwapChainBitDepth {
+        _8 = 0,
+        _10,
+        _16,
+    };
+    struct DisplayChromaticities {
+        float RedX{0.f};
+        float RedY{0.f};
+        float GreenX{0.f};
+        float GreenY{0.f};
+        float BlueX{0.f};
+        float BlueY{0.f};
+        float WhiteX{0.f};
+        float WhiteY{0.f};
+    };
+    struct DXSwapchainOption {
+        uint64_t window;
+        uint2 size;
+        PixelStorage storage = PixelStorage::HALF4;
+        bool wants_vsync = true;
+        uint back_buffer_count = 2;
+    };
+
+    enum class ColorSpace : uint {
+        RGB_FULL_G22_NONE_P709 = 0,
+        RGB_FULL_G10_NONE_P709 = 1,
+        RGB_STUDIO_G22_NONE_P709 = 2,
+        RGB_STUDIO_G22_NONE_P2020 = 3,
+        RESERVED = 4,
+        YCBCR_FULL_G22_NONE_P709_X601 = 5,
+        YCBCR_STUDIO_G22_LEFT_P601 = 6,
+        YCBCR_FULL_G22_LEFT_P601 = 7,
+        YCBCR_STUDIO_G22_LEFT_P709 = 8,
+        YCBCR_FULL_G22_LEFT_P709 = 9,
+        YCBCR_STUDIO_G22_LEFT_P2020 = 10,
+        YCBCR_FULL_G22_LEFT_P2020 = 11,
+        RGB_FULL_G2084_NONE_P2020 = 12,
+        YCBCR_STUDIO_G2084_LEFT_P2020 = 13,
+        RGB_STUDIO_G2084_NONE_P2020 = 14,
+        YCBCR_STUDIO_G22_TOPLEFT_P2020 = 15,
+        YCBCR_STUDIO_G2084_TOPLEFT_P2020 = 16,
+        RGB_FULL_G22_NONE_P2020 = 17,
+        YCBCR_STUDIO_GHLG_TOPLEFT_P2020 = 18,
+        YCBCR_FULL_GHLG_TOPLEFT_P2020 = 19,
+        RGB_STUDIO_G24_NONE_P709 = 20,
+        RGB_STUDIO_G24_NONE_P2020 = 21,
+        YCBCR_STUDIO_G24_LEFT_P709 = 22,
+        YCBCR_STUDIO_G24_LEFT_P2020 = 23,
+        YCBCR_STUDIO_G24_TOPLEFT_P2020 = 24,
+        CUSTOM = 0xFFFFFFFF
+    };
+
+    struct Meta {
+        ColorSpace color_space;
+        DisplayChromaticities chromaticities;
+    };
+
+    [[nodiscard]] virtual SwapchainCreationInfo create_swapchain(
+        const DXSwapchainOption &option,
+        uint64_t stream_handle) noexcept = 0;
+    virtual Meta set_hdr_meta_data(
+        uint64_t swapchain_handle,
+        float max_output_nits = 1000.0f,
+        float min_output_nits = 0.001f,
+        float max_cll = 2000.0f,
+        float max_fall = 500.0f,
+        const DisplayChromaticities *custom_chroma = nullptr) noexcept = 0;
+    static constexpr luisa::string_view name = "DXHDRExt";
+    [[nodiscard]] Swapchain create_swapchain(const Stream &stream, const DXSwapchainOption &option) noexcept;
+
+protected:
+    ~DXHDRExt() = default;
+};
+}// namespace luisa::compute

include/luisa/runtime/swapchain.h

@@ -23,6 +23,7 @@ class LC_RUNTIME_API Swapchain final : public Resource {
 private:
     friend class Device;
     friend class ResourceGenerator;
+    friend class DXHDRExt;
     PixelStorage _storage{};
     Swapchain(DeviceInterface *device, const SwapchainCreationInfo &create_info) noexcept;
     Swapchain(DeviceInterface *device, const SwapchainOption &option, uint64_t stream_handle) noexcept;

src/backends/common/hlsl/hlsl_codegen_util.cpp

@@ -1061,20 +1061,20 @@ void CodegenUtility::GetFunctionName(CallExpr const *expr, vstd::StringBuilder &
             str << "GroupMemoryBarrierWithGroupSync()"sv;
             return;
         case CallOp::RASTER_DISCARD:
-            LUISA_ASSERT(opt->funcType == CodegenStackData::FuncType::Pixel, "Raster-Discard can only be used in pixel shader");
+            LUISA_ASSERT(opt->isPixelShader, "Raster-Discard can only be used in pixel shader");
             str << "discard";
             return;
         case CallOp::DDX: {
             if (opt->isRaster) {
-                LUISA_ASSERT(opt->funcType == CodegenStackData::FuncType::Pixel, "ddx can only be used in pixel shader");
+                LUISA_ASSERT(opt->isPixelShader, "ddx can only be used in pixel shader");
                 str << "ddx"sv;
             } else {
                 str << "_ddx"sv;
             }
         } break;
         case CallOp::DDY: {
             if (opt->isRaster) {
-                LUISA_ASSERT(opt->funcType == CodegenStackData::FuncType::Pixel, "ddy can only be used in pixel shader");
+                LUISA_ASSERT(opt->isPixelShader, "ddy can only be used in pixel shader");
                 str << "ddy"sv;
             } else {
                 str << "_ddy"sv;
@@ -1580,20 +1580,7 @@ void main(uint3 thdId:SV_GroupThreadId,uint3 dspId:SV_DispatchThreadID,uint3 grp
     callable(callable, func);
 }
 void CodegenUtility::CodegenVertex(Function vert, vstd::StringBuilder &result, bool cBufferNonEmpty) {
-    vstd::unordered_set<void const *> callableMap;
-    auto gen = [&](auto &callable, Function func) -> void {
-        for (auto &&i : func.custom_callables()) {
-            if (callableMap.emplace(i.get()).second) {
-                Function f(i.get());
-                callable(callable, f);
-            }
-        }
-    };
-    auto callable = [&](auto &callable, Function func) -> void {
-        gen(callable, func);
-        CodegenFunction(func, result, cBufferNonEmpty);
-    };
-    gen(callable, vert);
+    CodegenFunction(vert, result, cBufferNonEmpty);
     auto args = vert.arguments();
     vstd::StringBuilder retName;
     auto retType = vert.return_type();
@@ -1630,20 +1617,7 @@ void CodegenUtility::CodegenVertex(Function vert, vstd::StringBuilder &result, b
 void CodegenUtility::CodegenPixel(Function pixel, vstd::StringBuilder &result, bool cBufferNonEmpty) {
     opt->isPixelShader = true;
     auto resetPixelShaderKey = vstd::scope_exit([&] { opt->isPixelShader = false; });
-    vstd::unordered_set<void const *> callableMap;
-    auto gen = [&](auto &callable, Function func) -> void {
-        for (auto &&i : func.custom_callables()) {
-            if (callableMap.emplace(i.get()).second) {
-                Function f(i.get());
-                callable(callable, f);
-            }
-        }
-    };
-    auto callable = [&](auto &callable, Function func) -> void {
-        gen(callable, func);
-        CodegenFunction(func, result, cBufferNonEmpty);
-    };
-    gen(callable, pixel);
+    CodegenFunction(pixel, result, cBufferNonEmpty);
     vstd::StringBuilder retName;
     auto retType = pixel.return_type();
     GetTypeName(*retType, retName, Usage::READ);

src/backends/dx/DXApi/dx_hdr_ext.cpp

@@ -41,7 +41,8 @@ DXHDRExt::DisplayCurve EnsureSwapChainColorSpace(
     }
     return result;
 }
-void SetHDRMetaData(
+DXHDRExt::DisplayChromaticities SetHDRMetaData(
+    DXGI_COLOR_SPACE_TYPE &colorSpace,
     LCSwapChain *swapchain,
     bool hdr_support,
     float MaxOutputNits /*=1000.0f*/,
@@ -50,13 +51,13 @@ void SetHDRMetaData(
     float MaxFALL /*=500.0f*/,
     const DXHDRExt::DisplayChromaticities *Chroma) {
     if (!swapchain) {
-        return;
+        return {};
     }
 
     // Clean the hdr metadata if the display doesn't support HDR
     if (!hdr_support) {
         ThrowIfFailed(swapchain->swapChain->SetHDRMetaData(DXGI_HDR_METADATA_TYPE_NONE, 0, nullptr));
-        return;
+        return {};
     }
 
     static const DXHDRExt::DisplayChromaticities DisplayChromaticityList[] =
@@ -66,7 +67,6 @@ void SetHDRMetaData(
         };
 
     // Select the chromaticity based on HDR format of the DWM.
-    DXGI_COLOR_SPACE_TYPE colorSpace = DXGI_COLOR_SPACE_CUSTOM;
     DXHDRExt::SwapChainBitDepth hitDepth;
     switch (swapchain->format) {
         case DXGI_FORMAT_R10G10B10A2_TYPELESS:
@@ -97,7 +97,7 @@ void SetHDRMetaData(
     } else {
         // Reset the metadata since this is not a supported HDR format.
         ThrowIfFailed(swapchain->swapChain->SetHDRMetaData(DXGI_HDR_METADATA_TYPE_NONE, 0, nullptr));
-        return;
+        return {};
     }
 
     // Set HDR meta data
@@ -117,6 +117,7 @@ void SetHDRMetaData(
     HDR10MetaData.MaxContentLightLevel = static_cast<UINT16>(MaxCLL);
     HDR10MetaData.MaxFrameAverageLightLevel = static_cast<UINT16>(MaxFALL);
     ThrowIfFailed(swapchain->swapChain->SetHDRMetaData(DXGI_HDR_METADATA_TYPE_HDR10, sizeof(DXGI_HDR_METADATA_HDR10), &HDR10MetaData));
+    return *Chroma;
 }
 }// namespace dx_hdr_ext_detail
 DXHDRExtImpl::DXHDRExtImpl(LCDevice *lc_device) : _lc_device(lc_device) {}
@@ -146,20 +147,25 @@ SwapchainCreationInfo DXHDRExtImpl::create_swapchain(
     return info;
 }
 
-void DXHDRExtImpl::set_hdr_meta_data(
+auto DXHDRExtImpl::set_hdr_meta_data(
     uint64_t swapchain_handle,
     float max_output_nits,
     float min_output_nits,
     float max_cll,
     float max_fall,
-    const DXHDRExt::DisplayChromaticities *custom_chroma) noexcept {
-    dx_hdr_ext_detail::SetHDRMetaData(
+    const DXHDRExt::DisplayChromaticities *custom_chroma) noexcept -> Meta {
+    DXGI_COLOR_SPACE_TYPE color_space = DXGI_COLOR_SPACE_CUSTOM;
+    auto chroma = dx_hdr_ext_detail::SetHDRMetaData(
+        color_space,
         reinterpret_cast<LCSwapChain *>(swapchain_handle),
         true,
         max_output_nits,
         min_output_nits,
         max_cll,
         max_fall,
         custom_chroma);
+    return {
+        static_cast<ColorSpace>(color_space),
+        chroma};
 }
 }// namespace lc::dx

src/backends/dx/DXApi/dx_hdr_ext.hpp

@@ -1,5 +1,5 @@
 #pragma once
-#include <luisa/backends/ext/dx_hdr_ext.h>
+#include <luisa/backends/ext/dx_hdr_ext_interface.h>
 #include <DXApi/LCDevice.h>
 namespace lc::dx {
 class DXHDRExtImpl : public luisa::compute::DXHDRExt {
@@ -10,7 +10,7 @@ class DXHDRExtImpl : public luisa::compute::DXHDRExt {
     SwapchainCreationInfo create_swapchain(
         const DXSwapchainOption &option,
         uint64_t stream_handle) noexcept override;
-    void set_hdr_meta_data(
+    Meta set_hdr_meta_data(
         uint64_t swapchain_handle,
         float max_output_nits = 1000.0f,
         float min_output_nits = 0.001f,

src/backends/dx/Shader/RasterShader.cpp

@@ -7,7 +7,14 @@
 #include <luisa/core/logging.h>
 namespace lc::dx {
 namespace RasterShaderDetail {
-static constexpr bool PRINT_CODE = false;
+static const bool PRINT_CODE = ([] {
+    // read env LUISA_DUMP_SOURCE
+    auto env = std::getenv("LUISA_DUMP_SOURCE");
+    if (env == nullptr) {
+        return false;
+    }
+    return std::string_view{env} == "1";
+})();
 static vstd::vector<SavedArgument> GetKernelArgs(Function vertexKernel, Function pixelKernel) {
     if (vertexKernel.builder() == nullptr || pixelKernel.builder() == nullptr) {
         return {};
@@ -278,7 +285,7 @@ RasterShader *RasterShader::CompileRaster(
         if (str.useBufferBindless) bdlsBufferCount++;
         if (str.useTex2DBindless) bdlsBufferCount++;
         if (str.useTex3DBindless) bdlsBufferCount++;
-        if constexpr (RasterShaderDetail::PRINT_CODE) {
+        if (RasterShaderDetail::PRINT_CODE) {
             auto f = fopen("hlsl_output.hlsl", "wb");
             fwrite(str.result.data(), str.result.size(), 1, f);
             fclose(f);
@@ -349,7 +356,7 @@ void RasterShader::SaveRaster(
     Function pixelKernel,
     uint shaderModel,
     bool enableUnsafeMath) {
-    if constexpr (RasterShaderDetail::PRINT_CODE) {
+    if (RasterShaderDetail::PRINT_CODE) {
         auto f = fopen("hlsl_output.hlsl", "ab");
         fwrite(str.result.data(), str.result.size(), 1, f);
         fclose(f);