Overhaul IR lowering of pointer types. (#4710)

* Overhaul IR lowering of pointer types.

* Propagate address space in IRBuilder.

* Fixup.

* Fix.

* Fix.

* Change how Ptr type is printed to text.

* Fix.

1 files changed, 46 insertions, 8 deletions

diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index ba7376c46..bb05a1c29 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -2813,9 +2813,9 @@ namespace Slang
         return (IRInOutType*) getPtrType(kIROp_InOutType, valueType);
     }
 
-    IRRefType* IRBuilder::getRefType(IRType* valueType)
+    IRRefType* IRBuilder::getRefType(IRType* valueType, AddressSpace addrSpace)
     {
-        return (IRRefType*) getPtrType(kIROp_RefType, valueType);
+        return (IRRefType*) getPtrType(kIROp_RefType, valueType, addrSpace);
     }
 
     IRConstRefType* IRBuilder::getConstRefType(IRType* valueType)
@@ -2840,8 +2840,13 @@ namespace Slang
 
     IRPtrType* IRBuilder::getPtrType(IROp op, IRType* valueType, IRIntegerValue addressSpace)
     {
-        IRInst* operands[] = {valueType, getIntValue(getIntType(), addressSpace)};
-        return (IRPtrType*)getType(op, 2, operands);
+        return (IRPtrType*)getPtrType(op, valueType, getIntValue(getUInt64Type(), addressSpace));
+    }
+
+    IRPtrType* IRBuilder::getPtrType(IROp op, IRType* valueType, IRInst* addressSpace)
+    {
+        IRInst* operands[] = { valueType, addressSpace };
+        return (IRPtrType*)getType(op, addressSpace ? 2 : 1, operands);
     }
 
     IRTextureTypeBase* IRBuilder::getTextureType(IRType* elementType, IRInst* shape, IRInst* isArray, IRInst* isMS, IRInst* sampleCount, IRInst* access, IRInst* isShadow, IRInst* isCombined, IRInst* format)
@@ -4881,11 +4886,30 @@ namespace Slang
         return inst;
     }
 
+    IRType* maybePropagateAddressSpace(IRBuilder* builder, IRInst* basePtr, IRType* type)
+    {
+        if (auto basePtrType = as<IRPtrTypeBase>(basePtr->getDataType()))
+        {
+            if (auto resultPtrType = as<IRPtrTypeBase>(type))
+            {
+                if (basePtrType->getAddressSpace() != resultPtrType->getAddressSpace())
+                {
+                    type = builder->getPtrType(
+                        resultPtrType->getOp(), resultPtrType->getValueType(), basePtrType->getAddressSpace());
+                }
+            }
+        }
+        return type;
+    }
+
     IRInst* IRBuilder::emitFieldAddress(
         IRType* type,
         IRInst* base,
         IRInst* field)
     {
+        // Propagate pointer address space if it is available on base.
+        type = maybePropagateAddressSpace(this, base, type);
+
         auto inst = createInst<IRFieldAddress>(
             this,
             kIROp_FieldAddress,
@@ -4982,6 +5006,9 @@ namespace Slang
         IRInst*    basePtr,
         IRInst*    index)
     {
+        // Propagate pointer address space if it is available on base.
+        type = maybePropagateAddressSpace(this, basePtr, type);
+
         auto inst = createInst<IRFieldAddress>(
             this,
             kIROp_GetElementPtr,
@@ -5004,9 +5031,20 @@ namespace Slang
         IRInst* basePtr,
         IRInst* index)
     {
+        AddressSpace addrSpace = AddressSpace::Generic;
+        IRInst* valueType = nullptr;
+        auto basePtrType = unwrapAttributedType(basePtr->getDataType());
+        if (auto ptrType = as<IRPtrTypeBase>(basePtrType))
+        {
+            addrSpace = ptrType->getAddressSpace();
+            valueType = ptrType->getValueType();
+        }
+        else if (auto ptrLikeType = as<IRPointerLikeType>(basePtrType))
+        {
+            valueType = ptrLikeType->getElementType();
+        }
         IRType* type = nullptr;
-        auto basePtrType = as<IRPtrTypeBase>(basePtr->getDataType());
-        auto valueType = unwrapAttributedType(basePtrType->getValueType());
+        valueType = unwrapAttributedType(valueType);
         if (auto arrayType = as<IRArrayTypeBase>(valueType))
         {
             type = arrayType->getElementType();
@@ -5028,7 +5066,7 @@ namespace Slang
         auto inst = createInst<IRGetElementPtr>(
             this,
             kIROp_GetElementPtr,
-            getPtrType(type),
+            getPtrType(kIROp_PtrType, type, addrSpace),
             basePtr,
             index);
 
@@ -5058,7 +5096,7 @@ namespace Slang
                     }
                 }
                 SLANG_RELEASE_ASSERT(resultType);
-                basePtr = emitFieldAddress(getPtrType(resultType), basePtr, structKey);
+                basePtr = emitFieldAddress(getPtrType(kIROp_PtrType, resultType, basePtrType->getAddressSpace()), basePtr, structKey);
             }
             else
             {