From 549aa897bcfedc28fb2ef8009396f846ea182b72 Mon Sep 17 00:00:00 2001
From: Mukund Keshava <mkeshava@nvidia.com>
Date: Tue, 1 Apr 2025 18:16:17 +0530
Subject: Add GetDimensions support for CUDA (#6718)

* Add GetDimensions support for CUDA

This CL adds GetDimensions support for cuda by using the PTX
instructions. Currently, PTX only supports getting width, height and
depth.

This CL also adds a new test to test this support.

Fixes #5139

* format code

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
---
 source/slang/slang-core-module-textures.cpp | 43 ++++++++++++++++++++++++++++-
 1 file changed, 42 insertions(+), 1 deletion(-)

(limited to 'source')

diff --git a/source/slang/slang-core-module-textures.cpp b/source/slang/slang-core-module-textures.cpp
index f703a8a3b..ed34ed388 100644
--- a/source/slang/slang-core-module-textures.cpp
+++ b/source/slang/slang-core-module-textures.cpp
@@ -258,6 +258,9 @@ void TextureTypeInfo::writeGetDimensionFunctions()
             StringBuilder metal;
             const char* metalMipLevel = "0";
 
+            StringBuilder cuda;
+            cuda << "{";
+
             StringBuilder wgsl;
             wgsl << "{";
 
@@ -277,6 +280,9 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 params << t << "width";
                 metal << "(*($" << String(paramCount) << ") = $0.get_width("
                       << String(metalMipLevel) << ")),";
+                cuda << "uint32_t width; asm(\\\"txq.width.b32 %0, [%1];\\\" : \\\"=r\\\"(width) : "
+                        "\\\"l\\\"($0)); *($"
+                     << String(paramCount) << ") = width;";
                 wgsl << "($" << String(paramCount) << ") = "
                      << wgslTextureAttributeConversion(
                             dimType,
@@ -292,6 +298,9 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 params << t << "width,";
                 metal << "(*($" << String(paramCount) << ") = $0.get_width("
                       << String(metalMipLevel) << ")),";
+                cuda << "uint32_t w, h; asm(\\\"txq.width.b32 %0, [%2]; txq.height.b32 %1, "
+                        "[%2];\\\" : \\\"=r\\\"(w), \\\"=r\\\"(h) : \\\"l\\\"($0)); *($"
+                     << String(paramCount) << ") = w;";
                 wgsl << "var dim = textureDimensions($0" << (includeMipInfo ? ", $1" : "") << ");";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "dim.x") << ";";
@@ -300,6 +309,9 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 params << t << "height";
                 metal << "(*($" << String(paramCount) << ") = $0.get_height("
                       << String(metalMipLevel) << ")),";
+                if (cuda.getLength() > 1 && cuda[cuda.getLength() - 1] != ';')
+                    cuda << "; ";
+                cuda << "*($" << String(paramCount) << ") = h;";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "dim.y") << ";";
 
@@ -311,6 +323,10 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 params << t << "width,";
                 metal << "(*($" << String(paramCount) << ") = $0.get_width("
                       << String(metalMipLevel) << ")),";
+                cuda << "uint32_t w, h, d; asm(\\\"txq.width.b32 %0, [%3]; txq.height.b32 %1, "
+                        "[%3]; txq.depth.b32 %2, [%3];\\\" : \\\"=r\\\"(w), \\\"=r\\\"(h), "
+                        "\\\"=r\\\"(d) : \\\"l\\\"($0)); *($"
+                     << String(paramCount) << ") = w;";
                 wgsl << "var dim = textureDimensions($0" << (includeMipInfo ? ", $1" : "") << ");";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "dim.x") << ";";
@@ -319,6 +335,9 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 params << t << "height,";
                 metal << "(*($" << String(paramCount) << ") = $0.get_height("
                       << String(metalMipLevel) << ")),";
+                if (cuda.getLength() > 1 && cuda[cuda.getLength() - 1] != ';')
+                    cuda << "; ";
+                cuda << "*($" << String(paramCount) << ") = h;";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "dim.y") << ";";
 
@@ -326,6 +345,9 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 params << t << "depth";
                 metal << "(*($" << String(paramCount) << ") = $0.get_depth("
                       << String(metalMipLevel) << ")),";
+                if (cuda.getLength() > 1 && cuda[cuda.getLength() - 1] != ';')
+                    cuda << "; ";
+                cuda << "*($" << String(paramCount) << ") = d;";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "dim.z") << ";";
 
@@ -343,6 +365,14 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 ++paramCount;
                 params << ", " << t << "elements";
                 metal << "(*($" << String(paramCount) << ") = $0.get_array_size()),";
+
+                // For cube map arrays, CUDA should include all 6 faces in the array size count
+                // but we can't currently implement this as txq.array_size isn't supported
+                if (cuda.getLength() > 1 && cuda[cuda.getLength() - 1] != ';')
+                    cuda << "; ";
+                cuda << "/* txq.array_size not available in CUDA */ *($" << String(paramCount)
+                     << ") = 0;";
+
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "textureNumLayers($0)")
                      << ";";
@@ -353,6 +383,10 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 ++paramCount;
                 params << ", " << t << "sampleCount";
                 metal << "(*($" << String(paramCount) << ") = $0.get_num_samples()),";
+                if (cuda.getLength() > 1 && cuda[cuda.getLength() - 1] != ';')
+                    cuda << "; ";
+                cuda << "/* txq.samples not available in CUDA */ *($" << String(paramCount)
+                     << ") = 0;";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "textureNumSamples($0)")
                      << ";";
@@ -363,12 +397,17 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 ++paramCount;
                 params << ", " << t << "numberOfLevels";
                 metal << "(*($" << String(paramCount) << ") = $0.get_num_mip_levels()),";
+                if (cuda.getLength() > 1 && cuda[cuda.getLength() - 1] != ';')
+                    cuda << "; ";
+                cuda << "/* txq.num_mipmap_levels not available in CUDA */ *($"
+                     << String(paramCount) << ") = 0;";
                 wgsl << "($" << String(paramCount)
                      << ") = " << wgslTextureAttributeConversion(dimType, "textureNumLevels($0)")
                      << ";";
             }
 
             metal.reduceLength(metal.getLength() - 1); // drop the last comma
+            cuda << "}";
             wgsl << "}";
 
             StringBuilder glsl;
@@ -559,6 +598,8 @@ void TextureTypeInfo::writeGetDimensionFunctions()
             sb << "    __glsl_version(450)\n";
 
             sb << "    [require(cpp";
+            if (cuda.getLength())
+                sb << "_cuda";
             if (glsl.getLength())
                 sb << "_glsl";
             sb << "_hlsl";
@@ -578,7 +619,7 @@ void TextureTypeInfo::writeGetDimensionFunctions()
                 spirvDefault,
                 spirvRWDefault,
                 spirvCombined,
-                "",
+                cuda.produceString(),
                 metal,
                 wgsl,
                 ReadNoneMode::Always);
-- 
cgit v1.2.3