From cc412af89e54b04ead508ca84825a18d001b92d0 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Thu, 31 Aug 2023 13:49:40 -0700
Subject: Add SPIRV atomics intrinsics and fix buffer layout lowering. (#3170)

* Fix atomics intrinsics and buffer layout lowering.

* Fix.

* Add more test.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 .../slang/slang-ir-lower-buffer-element-type.cpp   | 41 +++++++++++++++-------
 1 file changed, 29 insertions(+), 12 deletions(-)

(limited to 'source/slang/slang-ir-lower-buffer-element-type.cpp')

diff --git a/source/slang/slang-ir-lower-buffer-element-type.cpp b/source/slang/slang-ir-lower-buffer-element-type.cpp
index bf87d72fe..bc8e47e91 100644
--- a/source/slang/slang-ir-lower-buffer-element-type.cpp
+++ b/source/slang/slang-ir-lower-buffer-element-type.cpp
@@ -165,7 +165,9 @@ namespace Slang
             for (IRIntegerValue ii = 0; ii < count; ++ii)
             {
                 auto packedElement = builder.emitElementExtract(packedArray, ii);
-                auto originalElement = builder.emitCallInst(innerTypeInfo.originalType, innerTypeInfo.convertLoweredToOriginal, 1, &packedElement);
+                auto originalElement = innerTypeInfo.convertLoweredToOriginal
+                    ? builder.emitCallInst(innerTypeInfo.originalType, innerTypeInfo.convertLoweredToOriginal, 1, &packedElement)
+                    : packedElement;
                 args[(Index)ii] = originalElement;
             }
             auto result = builder.emitMakeArray(arrayType, (UInt)args.getCount(), args.getBuffer());
@@ -194,7 +196,9 @@ namespace Slang
             for (IRIntegerValue ii = 0; ii < count; ++ii)
             {
                 auto originalElement = builder.emitElementExtract(originalParam, ii);
-                auto packedElement = builder.emitCallInst(innerTypeInfo.loweredType, innerTypeInfo.convertOriginalToLowered, 1, &originalElement);
+                auto packedElement = innerTypeInfo.convertOriginalToLowered
+                    ? builder.emitCallInst(innerTypeInfo.loweredType, innerTypeInfo.convertOriginalToLowered, 1, &originalElement)
+                    : originalElement;
                 args[(Index)ii] = packedElement;
             }
             auto packedArray = builder.emitMakeArray(innerArrayType, (UInt)args.getCount(), args.getBuffer());
@@ -259,7 +263,7 @@ namespace Slang
                 auto arrayType = builder.getArrayType(
                     vectorType,
                     isColMajor?matrixType->getColumnCount():matrixType->getRowCount(),
-                    builder.getIntValue(builder.getIntType(), elementSizeAlignment.size));
+                    builder.getIntValue(builder.getIntType(), elementSizeAlignment.getStride()));
                 builder.createStructField(loweredType, structKey, arrayType);
 
                 info.loweredType = loweredType;
@@ -272,10 +276,16 @@ namespace Slang
             else if (auto arrayType = as<IRArrayType>(type))
             {
                 auto loweredInnerTypeInfo = getLoweredTypeInfo(arrayType->getElementType(), rules);
-                if (!loweredInnerTypeInfo.convertLoweredToOriginal)
+                // For spirv backend, we always want to lower all array types, even if the element type
+                // comes out the same. This is because different layout rules may have different array
+                // stride requirements.
+                if (!target->shouldEmitSPIRVDirectly())
                 {
-                    info.loweredType = type;
-                    return info;
+                    if (!loweredInnerTypeInfo.convertLoweredToOriginal)
+                    {
+                        info.loweredType = type;
+                        return info;
+                    }
                 }
                 auto loweredType = builder.createStructType();
                 info.loweredType = loweredType;
@@ -287,12 +297,12 @@ namespace Slang
                 auto structKey = builder.createStructKey();
                 builder.addNameHintDecoration(structKey, UnownedStringSlice("data"));
                 IRSizeAndAlignment elementSizeAlignment;
-                getSizeAndAlignment(rules, loweredType, &elementSizeAlignment);
+                getSizeAndAlignment(rules, loweredInnerTypeInfo.loweredType, &elementSizeAlignment);
                 elementSizeAlignment = rules->alignCompositeElement(elementSizeAlignment);
                 auto innerArrayType = builder.getArrayType(
                     loweredInnerTypeInfo.loweredType,
                     arrayType->getElementCount(),
-                    builder.getIntValue(builder.getIntType(), elementSizeAlignment.size));
+                    builder.getIntValue(builder.getIntType(), elementSizeAlignment.getStride()));
                 builder.createStructField(loweredType, structKey, innerArrayType);
                 info.loweredInnerArrayType = innerArrayType;
                 info.loweredInnerStructKey = structKey;
@@ -312,12 +322,19 @@ namespace Slang
                     if (loweredFieldTypeInfo.convertLoweredToOriginal || rules->ruleName != IRTypeLayoutRuleName::Natural)
                         isTrivial = false;
                 }
-                if (isTrivial)
+
+                // For spirv backend, we always want to lower all array types, even if the element type
+                // comes out the same. This is because different layout rules may have different array
+                // stride requirements.
+                if (!target->shouldEmitSPIRVDirectly())
                 {
-                    info.loweredType = type;
-                    return info;
+                    // For non-spirv target, we skip lowering this type if all field types are unchanged.
+                    if (isTrivial)
+                    {
+                        info.loweredType = type;
+                        return info;
+                    }
                 }
-
                 auto loweredType = builder.createStructType();
                 StringBuilder nameSB;
                 getTypeNameHint(nameSB, type);
-- 
cgit v1.2.3