1  using Ryujinx.Common.Memory;
   2  using Ryujinx.Graphics.GAL;
   3  using Ryujinx.Graphics.Shader;
   4  using Silk.NET.Vulkan;
   5  using System;
   6  using System.Buffers;
   7  using System.Collections.Generic;
   8  using System.Runtime.CompilerServices;
   9  using System.Runtime.InteropServices;
  10  using CompareOp = Ryujinx.Graphics.GAL.CompareOp;
  11  using Format = Ryujinx.Graphics.GAL.Format;
  12  using SamplerCreateInfo = Ryujinx.Graphics.GAL.SamplerCreateInfo;
  13  
  14  namespace Ryujinx.Graphics.Vulkan
  15  {
  16      class DescriptorSetUpdater
  17      {
  18          private const ulong StorageBufferMaxMirrorable = 0x2000;
  19  
  20          private const int ArrayGrowthSize = 16;
  21  
  22          private record struct BufferRef
  23          {
  24              public Auto<DisposableBuffer> Buffer;
  25              public int Offset;
  26              public bool Write;
  27  
  28              public BufferRef(Auto<DisposableBuffer> buffer)
  29              {
  30                  Buffer = buffer;
  31                  Offset = 0;
  32                  Write = true;
  33              }
  34  
  35              public BufferRef(Auto<DisposableBuffer> buffer, ref BufferRange range)
  36              {
  37                  Buffer = buffer;
  38                  Offset = range.Offset;
  39                  Write = range.Write;
  40              }
  41          }
  42  
  43          private record struct TextureRef
  44          {
  45              public ShaderStage Stage;
  46              public TextureView View;
  47              public Auto<DisposableImageView> ImageView;
  48              public Auto<DisposableSampler> Sampler;
  49  
  50              public TextureRef(ShaderStage stage, TextureView view, Auto<DisposableImageView> imageView, Auto<DisposableSampler> sampler)
  51              {
  52                  Stage = stage;
  53                  View = view;
  54                  ImageView = imageView;
  55                  Sampler = sampler;
  56              }
  57          }
  58  
  59          private record struct ImageRef
  60          {
  61              public ShaderStage Stage;
  62              public TextureView View;
  63              public Auto<DisposableImageView> ImageView;
  64  
  65              public ImageRef(ShaderStage stage, TextureView view, Auto<DisposableImageView> imageView)
  66              {
  67                  Stage = stage;
  68                  View = view;
  69                  ImageView = imageView;
  70              }
  71          }
  72  
  73          private readonly record struct ArrayRef<T>(ShaderStage Stage, T Array);
  74  
  75          private readonly VulkanRenderer _gd;
  76          private readonly Device _device;
  77          private ShaderCollection _program;
  78  
  79          private readonly BufferRef[] _uniformBufferRefs;
  80          private readonly BufferRef[] _storageBufferRefs;
  81          private readonly TextureRef[] _textureRefs;
  82          private readonly ImageRef[] _imageRefs;
  83          private readonly TextureBuffer[] _bufferTextureRefs;
  84          private readonly TextureBuffer[] _bufferImageRefs;
  85  
  86          private ArrayRef<TextureArray>[] _textureArrayRefs;
  87          private ArrayRef<ImageArray>[] _imageArrayRefs;
  88  
  89          private ArrayRef<TextureArray>[] _textureArrayExtraRefs;
  90          private ArrayRef<ImageArray>[] _imageArrayExtraRefs;
  91  
  92          private readonly DescriptorBufferInfo[] _uniformBuffers;
  93          private readonly DescriptorBufferInfo[] _storageBuffers;
  94          private readonly DescriptorImageInfo[] _textures;
  95          private readonly DescriptorImageInfo[] _images;
  96          private readonly BufferView[] _bufferTextures;
  97          private readonly BufferView[] _bufferImages;
  98  
  99          private readonly DescriptorSetTemplateUpdater _templateUpdater;
 100  
 101          private BitMapStruct<Array2<long>> _uniformSet;
 102          private BitMapStruct<Array2<long>> _storageSet;
 103          private BitMapStruct<Array2<long>> _uniformMirrored;
 104          private BitMapStruct<Array2<long>> _storageMirrored;
 105          private readonly int[] _uniformSetPd;
 106          private int _pdSequence = 1;
 107  
 108          private bool _updateDescriptorCacheCbIndex;
 109  
 110          [Flags]
 111          private enum DirtyFlags
 112          {
 113              None = 0,
 114              Uniform = 1 << 0,
 115              Storage = 1 << 1,
 116              Texture = 1 << 2,
 117              Image = 1 << 3,
 118              All = Uniform | Storage | Texture | Image,
 119          }
 120  
 121          private DirtyFlags _dirty;
 122  
 123          private readonly BufferHolder _dummyBuffer;
 124          private readonly TextureView _dummyTexture;
 125          private readonly SamplerHolder _dummySampler;
 126  
 127          public List<TextureView> FeedbackLoopHazards { get; private set; }
 128  
 129          public DescriptorSetUpdater(VulkanRenderer gd, Device device)
 130          {
 131              _gd = gd;
 132              _device = device;
 133  
 134              // Some of the bindings counts needs to be multiplied by 2 because we have buffer and
 135              // regular textures/images interleaved on the same descriptor set.
 136  
 137              _uniformBufferRefs = new BufferRef[Constants.MaxUniformBufferBindings];
 138              _storageBufferRefs = new BufferRef[Constants.MaxStorageBufferBindings];
 139              _textureRefs = new TextureRef[Constants.MaxTextureBindings * 2];
 140              _imageRefs = new ImageRef[Constants.MaxImageBindings * 2];
 141              _bufferTextureRefs = new TextureBuffer[Constants.MaxTextureBindings * 2];
 142              _bufferImageRefs = new TextureBuffer[Constants.MaxImageBindings * 2];
 143  
 144              _textureArrayRefs = Array.Empty<ArrayRef<TextureArray>>();
 145              _imageArrayRefs = Array.Empty<ArrayRef<ImageArray>>();
 146  
 147              _textureArrayExtraRefs = Array.Empty<ArrayRef<TextureArray>>();
 148              _imageArrayExtraRefs = Array.Empty<ArrayRef<ImageArray>>();
 149  
 150              _uniformBuffers = new DescriptorBufferInfo[Constants.MaxUniformBufferBindings];
 151              _storageBuffers = new DescriptorBufferInfo[Constants.MaxStorageBufferBindings];
 152              _textures = new DescriptorImageInfo[Constants.MaxTexturesPerStage];
 153              _images = new DescriptorImageInfo[Constants.MaxImagesPerStage];
 154              _bufferTextures = new BufferView[Constants.MaxTexturesPerStage];
 155              _bufferImages = new BufferView[Constants.MaxImagesPerStage];
 156  
 157              _uniformSetPd = new int[Constants.MaxUniformBufferBindings];
 158  
 159              var initialImageInfo = new DescriptorImageInfo
 160              {
 161                  ImageLayout = ImageLayout.General,
 162              };
 163  
 164              _textures.AsSpan().Fill(initialImageInfo);
 165              _images.AsSpan().Fill(initialImageInfo);
 166  
 167              if (gd.Capabilities.SupportsNullDescriptors)
 168              {
 169                  // If null descriptors are supported, we can pass null as the handle.
 170                  _dummyBuffer = null;
 171              }
 172              else
 173              {
 174                  // If null descriptors are not supported, we need to pass the handle of a dummy buffer on unused bindings.
 175                  _dummyBuffer = gd.BufferManager.Create(gd, 0x10000, forConditionalRendering: false, baseType: BufferAllocationType.DeviceLocal);
 176              }
 177  
 178              _dummyTexture = gd.CreateTextureView(new TextureCreateInfo(
 179                  1,
 180                  1,
 181                  1,
 182                  1,
 183                  1,
 184                  1,
 185                  1,
 186                  4,
 187                  Format.R8G8B8A8Unorm,
 188                  DepthStencilMode.Depth,
 189                  Target.Texture2D,
 190                  SwizzleComponent.Red,
 191                  SwizzleComponent.Green,
 192                  SwizzleComponent.Blue,
 193                  SwizzleComponent.Alpha));
 194  
 195              _dummySampler = (SamplerHolder)gd.CreateSampler(new SamplerCreateInfo(
 196                  MinFilter.Nearest,
 197                  MagFilter.Nearest,
 198                  false,
 199                  AddressMode.Repeat,
 200                  AddressMode.Repeat,
 201                  AddressMode.Repeat,
 202                  CompareMode.None,
 203                  CompareOp.Always,
 204                  new ColorF(0, 0, 0, 0),
 205                  0,
 206                  0,
 207                  0,
 208                  1f));
 209  
 210              _templateUpdater = new();
 211          }
 212  
 213          public void Initialize(bool isMainPipeline)
 214          {
 215              MemoryOwner<byte> dummyTextureData = MemoryOwner<byte>.RentCleared(4);
 216              _dummyTexture.SetData(dummyTextureData);
 217  
 218              if (isMainPipeline)
 219              {
 220                  FeedbackLoopHazards = new();
 221              }
 222          }
 223  
 224          private static bool BindingOverlaps(ref DescriptorBufferInfo info, int bindingOffset, int offset, int size)
 225          {
 226              return offset < bindingOffset + (int)info.Range && (offset + size) > bindingOffset;
 227          }
 228  
 229          internal void Rebind(Auto<DisposableBuffer> buffer, int offset, int size)
 230          {
 231              if (_program == null)
 232              {
 233                  return;
 234              }
 235  
 236              // Check stage bindings
 237  
 238              _uniformMirrored.Union(_uniformSet).SignalSet((int binding, int count) =>
 239              {
 240                  for (int i = 0; i < count; i++)
 241                  {
 242                      ref BufferRef bufferRef = ref _uniformBufferRefs[binding];
 243                      if (bufferRef.Buffer == buffer)
 244                      {
 245                          ref DescriptorBufferInfo info = ref _uniformBuffers[binding];
 246                          int bindingOffset = bufferRef.Offset;
 247  
 248                          if (BindingOverlaps(ref info, bindingOffset, offset, size))
 249                          {
 250                              _uniformSet.Clear(binding);
 251                              _uniformSetPd[binding] = 0;
 252                              SignalDirty(DirtyFlags.Uniform);
 253                          }
 254                      }
 255  
 256                      binding++;
 257                  }
 258              });
 259  
 260              _storageMirrored.Union(_storageSet).SignalSet((int binding, int count) =>
 261              {
 262                  for (int i = 0; i < count; i++)
 263                  {
 264                      ref BufferRef bufferRef = ref _storageBufferRefs[binding];
 265                      if (bufferRef.Buffer == buffer)
 266                      {
 267                          ref DescriptorBufferInfo info = ref _storageBuffers[binding];
 268                          int bindingOffset = bufferRef.Offset;
 269  
 270                          if (BindingOverlaps(ref info, bindingOffset, offset, size))
 271                          {
 272                              _storageSet.Clear(binding);
 273                              SignalDirty(DirtyFlags.Storage);
 274                          }
 275                      }
 276  
 277                      binding++;
 278                  }
 279              });
 280          }
 281  
 282          public void InsertBindingBarriers(CommandBufferScoped cbs)
 283          {
 284              if ((FeedbackLoopHazards?.Count ?? 0) > 0)
 285              {
 286                  // Clear existing hazards - they will be rebuilt.
 287  
 288                  foreach (TextureView hazard in FeedbackLoopHazards)
 289                  {
 290                      hazard.DecrementHazardUses();
 291                  }
 292  
 293                  FeedbackLoopHazards.Clear();
 294              }
 295  
 296              foreach (ResourceBindingSegment segment in _program.BindingSegments[PipelineBase.TextureSetIndex])
 297              {
 298                  if (segment.Type == ResourceType.TextureAndSampler)
 299                  {
 300                      if (!segment.IsArray)
 301                      {
 302                          for (int i = 0; i < segment.Count; i++)
 303                          {
 304                              ref var texture = ref _textureRefs[segment.Binding + i];
 305                              texture.View?.PrepareForUsage(cbs, texture.Stage.ConvertToPipelineStageFlags(), FeedbackLoopHazards);
 306                          }
 307                      }
 308                      else
 309                      {
 310                          ref var arrayRef = ref _textureArrayRefs[segment.Binding];
 311                          PipelineStageFlags stageFlags = arrayRef.Stage.ConvertToPipelineStageFlags();
 312                          arrayRef.Array?.QueueWriteToReadBarriers(cbs, stageFlags);
 313                      }
 314                  }
 315              }
 316  
 317              foreach (ResourceBindingSegment segment in _program.BindingSegments[PipelineBase.ImageSetIndex])
 318              {
 319                  if (segment.Type == ResourceType.Image)
 320                  {
 321                      if (!segment.IsArray)
 322                      {
 323                          for (int i = 0; i < segment.Count; i++)
 324                          {
 325                              ref var image = ref _imageRefs[segment.Binding + i];
 326                              image.View?.PrepareForUsage(cbs, image.Stage.ConvertToPipelineStageFlags(), FeedbackLoopHazards);
 327                          }
 328                      }
 329                      else
 330                      {
 331                          ref var arrayRef = ref _imageArrayRefs[segment.Binding];
 332                          PipelineStageFlags stageFlags = arrayRef.Stage.ConvertToPipelineStageFlags();
 333                          arrayRef.Array?.QueueWriteToReadBarriers(cbs, stageFlags);
 334                      }
 335                  }
 336              }
 337  
 338              for (int setIndex = PipelineBase.DescriptorSetLayouts; setIndex < _program.BindingSegments.Length; setIndex++)
 339              {
 340                  var bindingSegments = _program.BindingSegments[setIndex];
 341  
 342                  if (bindingSegments.Length == 0)
 343                  {
 344                      continue;
 345                  }
 346  
 347                  ResourceBindingSegment segment = bindingSegments[0];
 348  
 349                  if (segment.IsArray)
 350                  {
 351                      if (segment.Type == ResourceType.Texture ||
 352                          segment.Type == ResourceType.Sampler ||
 353                          segment.Type == ResourceType.TextureAndSampler ||
 354                          segment.Type == ResourceType.BufferTexture)
 355                      {
 356                          ref var arrayRef = ref _textureArrayExtraRefs[setIndex - PipelineBase.DescriptorSetLayouts];
 357                          PipelineStageFlags stageFlags = arrayRef.Stage.ConvertToPipelineStageFlags();
 358                          arrayRef.Array?.QueueWriteToReadBarriers(cbs, stageFlags);
 359                      }
 360                      else if (segment.Type == ResourceType.Image || segment.Type == ResourceType.BufferImage)
 361                      {
 362                          ref var arrayRef = ref _imageArrayExtraRefs[setIndex - PipelineBase.DescriptorSetLayouts];
 363                          PipelineStageFlags stageFlags = arrayRef.Stage.ConvertToPipelineStageFlags();
 364                          arrayRef.Array?.QueueWriteToReadBarriers(cbs, stageFlags);
 365                      }
 366                  }
 367              }
 368          }
 369  
 370          public void AdvancePdSequence()
 371          {
 372              if (++_pdSequence == 0)
 373              {
 374                  _pdSequence = 1;
 375              }
 376          }
 377  
 378          public void SetProgram(CommandBufferScoped cbs, ShaderCollection program, bool isBound)
 379          {
 380              if (!program.HasSameLayout(_program))
 381              {
 382                  // When the pipeline layout changes, push descriptor bindings are invalidated.
 383  
 384                  AdvancePdSequence();
 385              }
 386  
 387              _program = program;
 388              _updateDescriptorCacheCbIndex = true;
 389              _dirty = DirtyFlags.All;
 390          }
 391  
 392          public void SetImage(CommandBufferScoped cbs, ShaderStage stage, int binding, ITexture image)
 393          {
 394              if (image is TextureBuffer imageBuffer)
 395              {
 396                  _bufferImageRefs[binding] = imageBuffer;
 397              }
 398              else if (image is TextureView view)
 399              {
 400                  ref ImageRef iRef = ref _imageRefs[binding];
 401  
 402                  iRef.View?.ClearUsage(FeedbackLoopHazards);
 403                  view?.PrepareForUsage(cbs, stage.ConvertToPipelineStageFlags(), FeedbackLoopHazards);
 404  
 405                  iRef = new(stage, view, view.GetIdentityImageView());
 406              }
 407              else
 408              {
 409                  _imageRefs[binding] = default;
 410                  _bufferImageRefs[binding] = null;
 411              }
 412  
 413              SignalDirty(DirtyFlags.Image);
 414          }
 415  
 416          public void SetImage(int binding, Auto<DisposableImageView> image)
 417          {
 418              _imageRefs[binding] = new(ShaderStage.Compute, null, image);
 419  
 420              SignalDirty(DirtyFlags.Image);
 421          }
 422  
 423          public void SetStorageBuffers(CommandBuffer commandBuffer, ReadOnlySpan<BufferAssignment> buffers)
 424          {
 425              for (int i = 0; i < buffers.Length; i++)
 426              {
 427                  var assignment = buffers[i];
 428                  var buffer = assignment.Range;
 429                  int index = assignment.Binding;
 430  
 431                  Auto<DisposableBuffer> vkBuffer = buffer.Handle == BufferHandle.Null
 432                      ? null
 433                      : _gd.BufferManager.GetBuffer(commandBuffer, buffer.Handle, buffer.Write, isSSBO: true);
 434  
 435                  ref BufferRef currentBufferRef = ref _storageBufferRefs[index];
 436  
 437                  DescriptorBufferInfo info = new()
 438                  {
 439                      Offset = (ulong)buffer.Offset,
 440                      Range = (ulong)buffer.Size,
 441                  };
 442  
 443                  var newRef = new BufferRef(vkBuffer, ref buffer);
 444  
 445                  ref DescriptorBufferInfo currentInfo = ref _storageBuffers[index];
 446  
 447                  if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
 448                  {
 449                      _storageSet.Clear(index);
 450  
 451                      currentInfo = info;
 452                      currentBufferRef = newRef;
 453                  }
 454              }
 455  
 456              SignalDirty(DirtyFlags.Storage);
 457          }
 458  
 459          public void SetStorageBuffers(CommandBuffer commandBuffer, int first, ReadOnlySpan<Auto<DisposableBuffer>> buffers)
 460          {
 461              for (int i = 0; i < buffers.Length; i++)
 462              {
 463                  var vkBuffer = buffers[i];
 464                  int index = first + i;
 465  
 466                  ref BufferRef currentBufferRef = ref _storageBufferRefs[index];
 467  
 468                  DescriptorBufferInfo info = new()
 469                  {
 470                      Offset = 0,
 471                      Range = Vk.WholeSize,
 472                  };
 473  
 474                  BufferRef newRef = new(vkBuffer);
 475  
 476                  ref DescriptorBufferInfo currentInfo = ref _storageBuffers[index];
 477  
 478                  if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
 479                  {
 480                      _storageSet.Clear(index);
 481  
 482                      currentInfo = info;
 483                      currentBufferRef = newRef;
 484                  }
 485              }
 486  
 487              SignalDirty(DirtyFlags.Storage);
 488          }
 489  
 490          public void SetTextureAndSampler(
 491              CommandBufferScoped cbs,
 492              ShaderStage stage,
 493              int binding,
 494              ITexture texture,
 495              ISampler sampler)
 496          {
 497              if (texture is TextureBuffer textureBuffer)
 498              {
 499                  _bufferTextureRefs[binding] = textureBuffer;
 500              }
 501              else if (texture is TextureView view)
 502              {
 503                  ref TextureRef iRef = ref _textureRefs[binding];
 504  
 505                  iRef.View?.ClearUsage(FeedbackLoopHazards);
 506                  view?.PrepareForUsage(cbs, stage.ConvertToPipelineStageFlags(), FeedbackLoopHazards);
 507  
 508                  iRef = new(stage, view, view.GetImageView(), ((SamplerHolder)sampler)?.GetSampler());
 509              }
 510              else
 511              {
 512                  _textureRefs[binding] = default;
 513                  _bufferTextureRefs[binding] = null;
 514              }
 515  
 516              SignalDirty(DirtyFlags.Texture);
 517          }
 518  
 519          public void SetTextureAndSamplerIdentitySwizzle(
 520              CommandBufferScoped cbs,
 521              ShaderStage stage,
 522              int binding,
 523              ITexture texture,
 524              ISampler sampler)
 525          {
 526              if (texture is TextureView view)
 527              {
 528                  view.Storage.QueueWriteToReadBarrier(cbs, AccessFlags.ShaderReadBit, stage.ConvertToPipelineStageFlags());
 529  
 530                  _textureRefs[binding] = new(stage, view, view.GetIdentityImageView(), ((SamplerHolder)sampler)?.GetSampler());
 531  
 532                  SignalDirty(DirtyFlags.Texture);
 533              }
 534              else
 535              {
 536                  SetTextureAndSampler(cbs, stage, binding, texture, sampler);
 537              }
 538          }
 539  
 540          public void SetTextureArray(CommandBufferScoped cbs, ShaderStage stage, int binding, ITextureArray array)
 541          {
 542              ref ArrayRef<TextureArray> arrayRef = ref GetArrayRef(ref _textureArrayRefs, binding, ArrayGrowthSize);
 543  
 544              if (arrayRef.Stage != stage || arrayRef.Array != array)
 545              {
 546                  arrayRef.Array?.DecrementBindCount();
 547  
 548                  if (array is TextureArray textureArray)
 549                  {
 550                      textureArray.IncrementBindCount();
 551                      textureArray.QueueWriteToReadBarriers(cbs, stage.ConvertToPipelineStageFlags());
 552                  }
 553  
 554                  arrayRef = new ArrayRef<TextureArray>(stage, array as TextureArray);
 555  
 556                  SignalDirty(DirtyFlags.Texture);
 557              }
 558          }
 559  
 560          public void SetTextureArraySeparate(CommandBufferScoped cbs, ShaderStage stage, int setIndex, ITextureArray array)
 561          {
 562              ref ArrayRef<TextureArray> arrayRef = ref GetArrayRef(ref _textureArrayExtraRefs, setIndex - PipelineBase.DescriptorSetLayouts);
 563  
 564              if (arrayRef.Stage != stage || arrayRef.Array != array)
 565              {
 566                  arrayRef.Array?.DecrementBindCount();
 567  
 568                  if (array is TextureArray textureArray)
 569                  {
 570                      textureArray.IncrementBindCount();
 571                      textureArray.QueueWriteToReadBarriers(cbs, stage.ConvertToPipelineStageFlags());
 572                  }
 573  
 574                  arrayRef = new ArrayRef<TextureArray>(stage, array as TextureArray);
 575  
 576                  SignalDirty(DirtyFlags.Texture);
 577              }
 578          }
 579  
 580          public void SetImageArray(CommandBufferScoped cbs, ShaderStage stage, int binding, IImageArray array)
 581          {
 582              ref ArrayRef<ImageArray> arrayRef = ref GetArrayRef(ref _imageArrayRefs, binding, ArrayGrowthSize);
 583  
 584              if (arrayRef.Stage != stage || arrayRef.Array != array)
 585              {
 586                  arrayRef.Array?.DecrementBindCount();
 587  
 588                  if (array is ImageArray imageArray)
 589                  {
 590                      imageArray.IncrementBindCount();
 591                      imageArray.QueueWriteToReadBarriers(cbs, stage.ConvertToPipelineStageFlags());
 592                  }
 593  
 594                  arrayRef = new ArrayRef<ImageArray>(stage, array as ImageArray);
 595  
 596                  SignalDirty(DirtyFlags.Image);
 597              }
 598          }
 599  
 600          public void SetImageArraySeparate(CommandBufferScoped cbs, ShaderStage stage, int setIndex, IImageArray array)
 601          {
 602              ref ArrayRef<ImageArray> arrayRef = ref GetArrayRef(ref _imageArrayExtraRefs, setIndex - PipelineBase.DescriptorSetLayouts);
 603  
 604              if (arrayRef.Stage != stage || arrayRef.Array != array)
 605              {
 606                  arrayRef.Array?.DecrementBindCount();
 607  
 608                  if (array is ImageArray imageArray)
 609                  {
 610                      imageArray.IncrementBindCount();
 611                      imageArray.QueueWriteToReadBarriers(cbs, stage.ConvertToPipelineStageFlags());
 612                  }
 613  
 614                  arrayRef = new ArrayRef<ImageArray>(stage, array as ImageArray);
 615  
 616                  SignalDirty(DirtyFlags.Image);
 617              }
 618          }
 619  
 620          private static ref ArrayRef<T> GetArrayRef<T>(ref ArrayRef<T>[] array, int index, int growthSize = 1)
 621          {
 622              ArgumentOutOfRangeException.ThrowIfNegative(index);
 623  
 624              if (array.Length <= index)
 625              {
 626                  Array.Resize(ref array, index + growthSize);
 627              }
 628  
 629              return ref array[index];
 630          }
 631  
 632          public void SetUniformBuffers(CommandBuffer commandBuffer, ReadOnlySpan<BufferAssignment> buffers)
 633          {
 634              for (int i = 0; i < buffers.Length; i++)
 635              {
 636                  var assignment = buffers[i];
 637                  var buffer = assignment.Range;
 638                  int index = assignment.Binding;
 639  
 640                  Auto<DisposableBuffer> vkBuffer = buffer.Handle == BufferHandle.Null
 641                      ? null
 642                      : _gd.BufferManager.GetBuffer(commandBuffer, buffer.Handle, false);
 643  
 644                  ref BufferRef currentBufferRef = ref _uniformBufferRefs[index];
 645  
 646                  DescriptorBufferInfo info = new()
 647                  {
 648                      Offset = (ulong)buffer.Offset,
 649                      Range = (ulong)buffer.Size,
 650                  };
 651  
 652                  BufferRef newRef = new(vkBuffer, ref buffer);
 653  
 654                  ref DescriptorBufferInfo currentInfo = ref _uniformBuffers[index];
 655  
 656                  if (!currentBufferRef.Equals(newRef) || currentInfo.Range != info.Range)
 657                  {
 658                      _uniformSet.Clear(index);
 659                      _uniformSetPd[index] = 0;
 660  
 661                      currentInfo = info;
 662                      currentBufferRef = newRef;
 663                  }
 664              }
 665  
 666              SignalDirty(DirtyFlags.Uniform);
 667          }
 668  
 669          private void SignalDirty(DirtyFlags flag)
 670          {
 671              _dirty |= flag;
 672          }
 673  
 674          public void UpdateAndBindDescriptorSets(CommandBufferScoped cbs, PipelineBindPoint pbp)
 675          {
 676              if ((_dirty & DirtyFlags.All) == 0)
 677              {
 678                  return;
 679              }
 680  
 681              var program = _program;
 682  
 683              if (_dirty.HasFlag(DirtyFlags.Uniform))
 684              {
 685                  if (program.UsePushDescriptors)
 686                  {
 687                      UpdateAndBindUniformBufferPd(cbs);
 688                  }
 689                  else
 690                  {
 691                      UpdateAndBind(cbs, program, PipelineBase.UniformSetIndex, pbp);
 692                  }
 693              }
 694  
 695              if (_dirty.HasFlag(DirtyFlags.Storage))
 696              {
 697                  UpdateAndBind(cbs, program, PipelineBase.StorageSetIndex, pbp);
 698              }
 699  
 700              if (_dirty.HasFlag(DirtyFlags.Texture))
 701              {
 702                  if (program.UpdateTexturesWithoutTemplate)
 703                  {
 704                      UpdateAndBindTexturesWithoutTemplate(cbs, program, pbp);
 705                  }
 706                  else
 707                  {
 708                      UpdateAndBind(cbs, program, PipelineBase.TextureSetIndex, pbp);
 709                  }
 710              }
 711  
 712              if (_dirty.HasFlag(DirtyFlags.Image))
 713              {
 714                  UpdateAndBind(cbs, program, PipelineBase.ImageSetIndex, pbp);
 715              }
 716  
 717              if (program.BindingSegments.Length > PipelineBase.DescriptorSetLayouts)
 718              {
 719                  // Program is using extra sets, we need to bind those too.
 720  
 721                  BindExtraSets(cbs, program, pbp);
 722              }
 723  
 724              _dirty = DirtyFlags.None;
 725          }
 726  
 727          [MethodImpl(MethodImplOptions.AggressiveInlining)]
 728          private static bool UpdateBuffer(
 729              CommandBufferScoped cbs,
 730              ref DescriptorBufferInfo info,
 731              ref BufferRef buffer,
 732              Auto<DisposableBuffer> dummyBuffer,
 733              bool mirrorable)
 734          {
 735              int offset = buffer.Offset;
 736              bool mirrored = false;
 737  
 738              if (mirrorable)
 739              {
 740                  info.Buffer = buffer.Buffer?.GetMirrorable(cbs, ref offset, (int)info.Range, out mirrored).Value ?? default;
 741              }
 742              else
 743              {
 744                  info.Buffer = buffer.Buffer?.Get(cbs, offset, (int)info.Range, buffer.Write).Value ?? default;
 745              }
 746  
 747              info.Offset = (ulong)offset;
 748  
 749              // The spec requires that buffers with null handle have offset as 0 and range as VK_WHOLE_SIZE.
 750              if (info.Buffer.Handle == 0)
 751              {
 752                  info.Buffer = dummyBuffer?.Get(cbs).Value ?? default;
 753                  info.Offset = 0;
 754                  info.Range = Vk.WholeSize;
 755              }
 756  
 757              return mirrored;
 758          }
 759  
 760          [MethodImpl(MethodImplOptions.AggressiveInlining)]
 761          private void UpdateAndBind(CommandBufferScoped cbs, ShaderCollection program, int setIndex, PipelineBindPoint pbp)
 762          {
 763              var bindingSegments = program.BindingSegments[setIndex];
 764  
 765              if (bindingSegments.Length == 0)
 766              {
 767                  return;
 768              }
 769  
 770              var dummyBuffer = _dummyBuffer?.GetBuffer();
 771  
 772              if (_updateDescriptorCacheCbIndex)
 773              {
 774                  _updateDescriptorCacheCbIndex = false;
 775                  program.UpdateDescriptorCacheCommandBufferIndex(cbs.CommandBufferIndex);
 776              }
 777  
 778              var dsc = program.GetNewDescriptorSetCollection(setIndex, out var isNew).Get(cbs);
 779  
 780              if (!program.HasMinimalLayout)
 781              {
 782                  if (isNew)
 783                  {
 784                      Initialize(cbs, setIndex, dsc);
 785                  }
 786              }
 787  
 788              DescriptorSetTemplate template = program.Templates[setIndex];
 789  
 790              DescriptorSetTemplateWriter tu = _templateUpdater.Begin(template);
 791  
 792              foreach (ResourceBindingSegment segment in bindingSegments)
 793              {
 794                  int binding = segment.Binding;
 795                  int count = segment.Count;
 796  
 797                  if (setIndex == PipelineBase.UniformSetIndex)
 798                  {
 799                      for (int i = 0; i < count; i++)
 800                      {
 801                          int index = binding + i;
 802  
 803                          if (_uniformSet.Set(index))
 804                          {
 805                              ref BufferRef buffer = ref _uniformBufferRefs[index];
 806  
 807                              bool mirrored = UpdateBuffer(cbs, ref _uniformBuffers[index], ref buffer, dummyBuffer, true);
 808  
 809                              _uniformMirrored.Set(index, mirrored);
 810                          }
 811                      }
 812  
 813                      ReadOnlySpan<DescriptorBufferInfo> uniformBuffers = _uniformBuffers;
 814  
 815                      tu.Push(uniformBuffers.Slice(binding, count));
 816                  }
 817                  else if (setIndex == PipelineBase.StorageSetIndex)
 818                  {
 819                      for (int i = 0; i < count; i++)
 820                      {
 821                          int index = binding + i;
 822  
 823                          ref BufferRef buffer = ref _storageBufferRefs[index];
 824  
 825                          if (_storageSet.Set(index))
 826                          {
 827                              ref var info = ref _storageBuffers[index];
 828  
 829                              bool mirrored = UpdateBuffer(cbs,
 830                                  ref info,
 831                                  ref _storageBufferRefs[index],
 832                                  dummyBuffer,
 833                                  !buffer.Write && info.Range <= StorageBufferMaxMirrorable);
 834  
 835                              _storageMirrored.Set(index, mirrored);
 836                          }
 837                      }
 838  
 839                      ReadOnlySpan<DescriptorBufferInfo> storageBuffers = _storageBuffers;
 840  
 841                      tu.Push(storageBuffers.Slice(binding, count));
 842                  }
 843                  else if (setIndex == PipelineBase.TextureSetIndex)
 844                  {
 845                      if (!segment.IsArray)
 846                      {
 847                          if (segment.Type != ResourceType.BufferTexture)
 848                          {
 849                              Span<DescriptorImageInfo> textures = _textures;
 850  
 851                              for (int i = 0; i < count; i++)
 852                              {
 853                                  ref var texture = ref textures[i];
 854                                  ref var refs = ref _textureRefs[binding + i];
 855  
 856                                  texture.ImageView = refs.ImageView?.Get(cbs).Value ?? default;
 857                                  texture.Sampler = refs.Sampler?.Get(cbs).Value ?? default;
 858  
 859                                  if (texture.ImageView.Handle == 0)
 860                                  {
 861                                      texture.ImageView = _dummyTexture.GetImageView().Get(cbs).Value;
 862                                  }
 863  
 864                                  if (texture.Sampler.Handle == 0)
 865                                  {
 866                                      texture.Sampler = _dummySampler.GetSampler().Get(cbs).Value;
 867                                  }
 868                              }
 869  
 870                              tu.Push<DescriptorImageInfo>(textures[..count]);
 871                          }
 872                          else
 873                          {
 874                              Span<BufferView> bufferTextures = _bufferTextures;
 875  
 876                              for (int i = 0; i < count; i++)
 877                              {
 878                                  bufferTextures[i] = _bufferTextureRefs[binding + i]?.GetBufferView(cbs, false) ?? default;
 879                              }
 880  
 881                              tu.Push<BufferView>(bufferTextures[..count]);
 882                          }
 883                      }
 884                      else
 885                      {
 886                          if (segment.Type != ResourceType.BufferTexture)
 887                          {
 888                              tu.Push(_textureArrayRefs[binding].Array.GetImageInfos(_gd, cbs, _dummyTexture, _dummySampler));
 889                          }
 890                          else
 891                          {
 892                              tu.Push(_textureArrayRefs[binding].Array.GetBufferViews(cbs));
 893                          }
 894                      }
 895                  }
 896                  else if (setIndex == PipelineBase.ImageSetIndex)
 897                  {
 898                      if (!segment.IsArray)
 899                      {
 900                          if (segment.Type != ResourceType.BufferImage)
 901                          {
 902                              Span<DescriptorImageInfo> images = _images;
 903  
 904                              for (int i = 0; i < count; i++)
 905                              {
 906                                  images[i].ImageView = _imageRefs[binding + i].ImageView?.Get(cbs).Value ?? default;
 907                              }
 908  
 909                              tu.Push<DescriptorImageInfo>(images[..count]);
 910                          }
 911                          else
 912                          {
 913                              Span<BufferView> bufferImages = _bufferImages;
 914  
 915                              for (int i = 0; i < count; i++)
 916                              {
 917                                  bufferImages[i] = _bufferImageRefs[binding + i]?.GetBufferView(cbs, true) ?? default;
 918                              }
 919  
 920                              tu.Push<BufferView>(bufferImages[..count]);
 921                          }
 922                      }
 923                      else
 924                      {
 925                          if (segment.Type != ResourceType.BufferTexture)
 926                          {
 927                              tu.Push(_imageArrayRefs[binding].Array.GetImageInfos(_gd, cbs, _dummyTexture));
 928                          }
 929                          else
 930                          {
 931                              tu.Push(_imageArrayRefs[binding].Array.GetBufferViews(cbs));
 932                          }
 933                      }
 934                  }
 935              }
 936  
 937              var sets = dsc.GetSets();
 938              _templateUpdater.Commit(_gd, _device, sets[0]);
 939  
 940              _gd.Api.CmdBindDescriptorSets(cbs.CommandBuffer, pbp, _program.PipelineLayout, (uint)setIndex, 1, sets, 0, ReadOnlySpan<uint>.Empty);
 941          }
 942  
 943          private void UpdateAndBindTexturesWithoutTemplate(CommandBufferScoped cbs, ShaderCollection program, PipelineBindPoint pbp)
 944          {
 945              int setIndex = PipelineBase.TextureSetIndex;
 946              var bindingSegments = program.BindingSegments[setIndex];
 947  
 948              if (bindingSegments.Length == 0)
 949              {
 950                  return;
 951              }
 952  
 953              if (_updateDescriptorCacheCbIndex)
 954              {
 955                  _updateDescriptorCacheCbIndex = false;
 956                  program.UpdateDescriptorCacheCommandBufferIndex(cbs.CommandBufferIndex);
 957              }
 958  
 959              var dsc = program.GetNewDescriptorSetCollection(setIndex, out _).Get(cbs);
 960  
 961              foreach (ResourceBindingSegment segment in bindingSegments)
 962              {
 963                  int binding = segment.Binding;
 964                  int count = segment.Count;
 965  
 966                  if (!segment.IsArray)
 967                  {
 968                      if (segment.Type != ResourceType.BufferTexture)
 969                      {
 970                          Span<DescriptorImageInfo> textures = _textures;
 971  
 972                          for (int i = 0; i < count; i++)
 973                          {
 974                              ref var texture = ref textures[i];
 975                              ref var refs = ref _textureRefs[binding + i];
 976  
 977                              texture.ImageView = refs.ImageView?.Get(cbs).Value ?? default;
 978                              texture.Sampler = refs.Sampler?.Get(cbs).Value ?? default;
 979  
 980                              if (texture.ImageView.Handle == 0)
 981                              {
 982                                  texture.ImageView = _dummyTexture.GetImageView().Get(cbs).Value;
 983                              }
 984  
 985                              if (texture.Sampler.Handle == 0)
 986                              {
 987                                  texture.Sampler = _dummySampler.GetSampler().Get(cbs).Value;
 988                              }
 989                          }
 990  
 991                          dsc.UpdateImages(0, binding, textures[..count], DescriptorType.CombinedImageSampler);
 992                      }
 993                      else
 994                      {
 995                          Span<BufferView> bufferTextures = _bufferTextures;
 996  
 997                          for (int i = 0; i < count; i++)
 998                          {
 999                              bufferTextures[i] = _bufferTextureRefs[binding + i]?.GetBufferView(cbs, false) ?? default;
1000                          }
1001  
1002                          dsc.UpdateBufferImages(0, binding, bufferTextures[..count], DescriptorType.UniformTexelBuffer);
1003                      }
1004                  }
1005                  else
1006                  {
1007                      if (segment.Type != ResourceType.BufferTexture)
1008                      {
1009                          dsc.UpdateImages(0, binding, _textureArrayRefs[binding].Array.GetImageInfos(_gd, cbs, _dummyTexture, _dummySampler), DescriptorType.CombinedImageSampler);
1010                      }
1011                      else
1012                      {
1013                          dsc.UpdateBufferImages(0, binding, _textureArrayRefs[binding].Array.GetBufferViews(cbs), DescriptorType.UniformTexelBuffer);
1014                      }
1015                  }
1016              }
1017  
1018              var sets = dsc.GetSets();
1019  
1020              _gd.Api.CmdBindDescriptorSets(cbs.CommandBuffer, pbp, _program.PipelineLayout, (uint)setIndex, 1, sets, 0, ReadOnlySpan<uint>.Empty);
1021          }
1022  
1023          [MethodImpl(MethodImplOptions.AggressiveInlining)]
1024          private void UpdateAndBindUniformBufferPd(CommandBufferScoped cbs)
1025          {
1026              int sequence = _pdSequence;
1027              var bindingSegments = _program.BindingSegments[PipelineBase.UniformSetIndex];
1028              var dummyBuffer = _dummyBuffer?.GetBuffer();
1029  
1030              long updatedBindings = 0;
1031              DescriptorSetTemplateWriter writer = _templateUpdater.Begin(32 * Unsafe.SizeOf<DescriptorBufferInfo>());
1032  
1033              foreach (ResourceBindingSegment segment in bindingSegments)
1034              {
1035                  int binding = segment.Binding;
1036                  int count = segment.Count;
1037  
1038                  ReadOnlySpan<DescriptorBufferInfo> uniformBuffers = _uniformBuffers;
1039  
1040                  for (int i = 0; i < count; i++)
1041                  {
1042                      int index = binding + i;
1043  
1044                      if (_uniformSet.Set(index))
1045                      {
1046                          ref BufferRef buffer = ref _uniformBufferRefs[index];
1047  
1048                          bool mirrored = UpdateBuffer(cbs, ref _uniformBuffers[index], ref buffer, dummyBuffer, true);
1049  
1050                          _uniformMirrored.Set(index, mirrored);
1051                      }
1052  
1053                      if (_uniformSetPd[index] != sequence)
1054                      {
1055                          // Need to set this push descriptor (even if the buffer binding has not changed)
1056  
1057                          _uniformSetPd[index] = sequence;
1058                          updatedBindings |= 1L << index;
1059  
1060                          writer.Push(MemoryMarshal.CreateReadOnlySpan(ref _uniformBuffers[index], 1));
1061                      }
1062                  }
1063              }
1064  
1065              if (updatedBindings > 0)
1066              {
1067                  DescriptorSetTemplate template = _program.GetPushDescriptorTemplate(updatedBindings);
1068                  _templateUpdater.CommitPushDescriptor(_gd, cbs, template, _program.PipelineLayout);
1069              }
1070          }
1071  
1072          private void Initialize(CommandBufferScoped cbs, int setIndex, DescriptorSetCollection dsc)
1073          {
1074              // We don't support clearing texture descriptors currently.
1075              if (setIndex != PipelineBase.UniformSetIndex && setIndex != PipelineBase.StorageSetIndex)
1076              {
1077                  return;
1078              }
1079  
1080              var dummyBuffer = _dummyBuffer?.GetBuffer().Get(cbs).Value ?? default;
1081  
1082              foreach (ResourceBindingSegment segment in _program.ClearSegments[setIndex])
1083              {
1084                  dsc.InitializeBuffers(0, segment.Binding, segment.Count, segment.Type.Convert(), dummyBuffer);
1085              }
1086          }
1087  
1088          private void BindExtraSets(CommandBufferScoped cbs, ShaderCollection program, PipelineBindPoint pbp)
1089          {
1090              for (int setIndex = PipelineBase.DescriptorSetLayouts; setIndex < program.BindingSegments.Length; setIndex++)
1091              {
1092                  var bindingSegments = program.BindingSegments[setIndex];
1093  
1094                  if (bindingSegments.Length == 0)
1095                  {
1096                      continue;
1097                  }
1098  
1099                  ResourceBindingSegment segment = bindingSegments[0];
1100  
1101                  if (segment.IsArray)
1102                  {
1103                      DescriptorSet[] sets = null;
1104  
1105                      if (segment.Type == ResourceType.Texture ||
1106                          segment.Type == ResourceType.Sampler ||
1107                          segment.Type == ResourceType.TextureAndSampler ||
1108                          segment.Type == ResourceType.BufferTexture)
1109                      {
1110                          sets = _textureArrayExtraRefs[setIndex - PipelineBase.DescriptorSetLayouts].Array.GetDescriptorSets(
1111                              _device,
1112                              cbs,
1113                              _templateUpdater,
1114                              program,
1115                              setIndex,
1116                              _dummyTexture,
1117                              _dummySampler);
1118                      }
1119                      else if (segment.Type == ResourceType.Image || segment.Type == ResourceType.BufferImage)
1120                      {
1121                          sets = _imageArrayExtraRefs[setIndex - PipelineBase.DescriptorSetLayouts].Array.GetDescriptorSets(
1122                              _device,
1123                              cbs,
1124                              _templateUpdater,
1125                              program,
1126                              setIndex,
1127                              _dummyTexture);
1128                      }
1129  
1130                      if (sets != null)
1131                      {
1132                          _gd.Api.CmdBindDescriptorSets(cbs.CommandBuffer, pbp, _program.PipelineLayout, (uint)setIndex, 1, sets, 0, ReadOnlySpan<uint>.Empty);
1133                      }
1134                  }
1135              }
1136          }
1137  
1138          public void SignalCommandBufferChange()
1139          {
1140              _updateDescriptorCacheCbIndex = true;
1141              _dirty = DirtyFlags.All;
1142  
1143              _uniformSet.Clear();
1144              _storageSet.Clear();
1145              AdvancePdSequence();
1146          }
1147  
1148          public void ForceTextureDirty()
1149          {
1150              SignalDirty(DirtyFlags.Texture);
1151          }
1152  
1153          public void ForceImageDirty()
1154          {
1155              SignalDirty(DirtyFlags.Image);
1156          }
1157  
1158          private static void SwapBuffer(BufferRef[] list, Auto<DisposableBuffer> from, Auto<DisposableBuffer> to)
1159          {
1160              for (int i = 0; i < list.Length; i++)
1161              {
1162                  if (list[i].Buffer == from)
1163                  {
1164                      list[i].Buffer = to;
1165                  }
1166              }
1167          }
1168  
1169          public void SwapBuffer(Auto<DisposableBuffer> from, Auto<DisposableBuffer> to)
1170          {
1171              SwapBuffer(_uniformBufferRefs, from, to);
1172              SwapBuffer(_storageBufferRefs, from, to);
1173          }
1174  
1175          protected virtual void Dispose(bool disposing)
1176          {
1177              if (disposing)
1178              {
1179                  _dummyTexture.Dispose();
1180                  _dummySampler.Dispose();
1181                  _templateUpdater.Dispose();
1182              }
1183          }
1184  
1185          public void Dispose()
1186          {
1187              Dispose(true);
1188          }
1189      }
1190  }