[SM6.10] Collect and Validate LinAlg PSV0 runtime data

lib/DxilValidation/DxilContainerValidation.cpp

@@ -137,7 +137,7 @@ class SemanticIndexTableVerifier {
       return false;
     if (Offset > Table.Entries)
       return false;
-    if ((Offset + Size) > Table.Entries)
+    if (Size > Table.Entries - Offset)
       return false;
     for (unsigned i = Offset; i < (Offset + Size); ++i) {
       UseMask[i] = true;
@@ -180,6 +180,7 @@ class PSVContentVerifier {
                               PSVSignatureElement0 *, const PSVStringTable &,
                               const PSVSemanticIndexTable &, std::string, bool);
   void VerifyResources(unsigned PSVVersion);
+  void VerifyLinAlgRuntimeInfo(unsigned PSVVersion);
   template <typename T>
   void VerifyResourceTable(T &ResTab, unsigned &ResourceIndex,
                            unsigned PSVVersion);
@@ -465,6 +466,190 @@ void PSVContentVerifier::VerifyEntryProperties(
   }
 }
 
+void PSVContentVerifier::VerifyLinAlgRuntimeInfo(unsigned PSVVersion) {
+  // Regenerate the expected runtime info to compare the container
+  // contents against
+  unique_ptr<DxilPartWriter> pWriter(NewPSVWriter(DM, PSVVersion));
+  CComPtr<AbstractMemoryStream> pOutputStream;
+  IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pOutputStream));
+  pOutputStream->Reserve(pWriter->size());
+  pWriter->write(pOutputStream);
+
+  DxilPipelineStateValidation ExpectedPSV;
+  if (!ExpectedPSV.InitFromPSV0(pOutputStream->GetPtr(),
+                                pOutputStream->GetPtrSize())) {
+    ValCtx.EmitFormatError(
+        ValidationRule::ContainerPartMatches,
+        {"Pipeline State Validation generated from DxilModule"});
+    return;
+  }
+
+  bool HasLinAlgRuntimeInfo = PSV.GetPSVLinAlgRuntimeInfo0() != nullptr;
+  bool ExpectedHasLinAlgRuntimeInfo =
+      ExpectedPSV.GetPSVLinAlgRuntimeInfo0() != nullptr;
+  if (HasLinAlgRuntimeInfo != ExpectedHasLinAlgRuntimeInfo) {
+    EmitMismatchError("LinAlgRuntimeInfoPresent",
+                      HasLinAlgRuntimeInfo ? "true" : "false",
+                      ExpectedHasLinAlgRuntimeInfo ? "true" : "false");
+    return;
+  }
+
+  if (!HasLinAlgRuntimeInfo)
+    return;
+
+  auto VerifyShapeReference =
+      [&](StringRef Name, const PSVLinAlgMatrixShapeArrayReference &ShapeRef,
+          const PSVLinAlgMatrixShapeArrayReference *ExpectedShapeRef) {
+        if (!IndexTableVerifier.MarkUse(ShapeRef.ShapesIndex, ShapeRef.Count)) {
+          EmitInvalidError("LinAlgOperationShapes");
+          return;
+        }
+        const uint32_t *ShapeIndexes =
+            PSV.GetSemanticIndexTable().Get(ShapeRef.ShapesIndex);
+        for (uint32_t I = 0; I < ShapeRef.Count; ++I) {
+          if (!PSV.GetPSVLinAlgMatrixOperationShape(ShapeIndexes[I])) {
+            EmitInvalidError("LinAlgOperationShapeIndex");
+            return;
+          }
+        }
+
+        if (!ExpectedShapeRef)
+          return;
+        if (ShapeRef.Count != ExpectedShapeRef->Count) {
+          EmitMismatchError((Name + "Count").str(),
+                            std::to_string(ShapeRef.Count),
+                            std::to_string(ExpectedShapeRef->Count));
+          return;
+        }
+        if (ShapeRef.Count == 0)
+          return;
+
+        const PSVSemanticIndexTable &ExpectedIndexTable =
+            ExpectedPSV.GetSemanticIndexTable();
+        if (ExpectedIndexTable.Table == nullptr ||
+            ExpectedShapeRef->ShapesIndex > ExpectedIndexTable.Entries ||
+            ExpectedShapeRef->Count >
+                ExpectedIndexTable.Entries - ExpectedShapeRef->ShapesIndex) {
+          EmitMismatchError(
+              Name, "valid shape index sequence",
+              "invalid shape index sequence generated from DxilModule");
+          return;
+        }
+
+        const uint32_t *ExpectedShapeIndexes =
+            ExpectedIndexTable.Get(ExpectedShapeRef->ShapesIndex);
+        if (!std::equal(ShapeIndexes, ShapeIndexes + ShapeRef.Count,
+                        ExpectedShapeIndexes))
+          EmitMismatchError(Name, "shape index sequence",
+                            "shape index sequence generated from DxilModule");
+      };
+
+  auto GetRecordName = [](StringRef Name, uint32_t I) {
+    return Name.str() + "[" + std::to_string(I) + "]";
+  };
+
+  auto GetRecordBytes = [](const auto &Record) {
+    static constexpr char HexDigits[] = "0123456789abcdef";
+    const uint8_t *Bytes = reinterpret_cast<const uint8_t *>(&Record);
+    std::string Result;
+    Result.reserve(sizeof(Record) * 3 - 1);
+    for (size_t I = 0; I < sizeof(Record); ++I) {
+      if (I != 0)
+        Result.push_back(' ');
+      Result.push_back(HexDigits[Bytes[I] >> 4]);
+      Result.push_back(HexDigits[Bytes[I] & 0xf]);
+    }
+    return Result;
+  };
+
+  auto VerifyRecord = [&](StringRef Name, uint32_t I, const auto &Record,
+                          const auto *ExpectedRecord) {
+    if (ExpectedRecord && memcmp(&Record, ExpectedRecord, sizeof(Record)) != 0)
+      EmitMismatchError(GetRecordName(Name, I), GetRecordBytes(Record),
+                        GetRecordBytes(*ExpectedRecord));
+  };
+
+  auto VerifyRecordWithShapes = [&](StringRef Name, uint32_t I,
+                                    const auto &Record,
+                                    const auto *ExpectedRecord) {
+    if (ExpectedRecord) {
+      auto ComparableRecord = Record;
+      auto ComparableExpectedRecord = *ExpectedRecord;
+      ComparableRecord.OperationShapes = {};
+      ComparableExpectedRecord.OperationShapes = {};
+      if (memcmp(&ComparableRecord, &ComparableExpectedRecord,
+                 sizeof(ComparableRecord)) != 0)
+        EmitMismatchError(GetRecordName(Name, I),
+                          GetRecordBytes(ComparableRecord),
+                          GetRecordBytes(ComparableExpectedRecord));
+    }
+
+    std::string ShapeName = Name.str() + "OperationShapes";
+    VerifyShapeReference(ShapeName, Record.OperationShapes,
+                         ExpectedRecord ? &ExpectedRecord->OperationShapes
+                                        : nullptr);
+  };
+
+  auto VerifyLinAlgTable = [&](StringRef Name, auto CountMethod, auto GetMethod,
+                               auto VerifyTableRecord) {
+    uint32_t Count = (PSV.*CountMethod)();
+    uint32_t ExpectedCount = (ExpectedPSV.*CountMethod)();
+    if (Count != ExpectedCount)
+      EmitMismatchError(Name.str() + "Count", std::to_string(Count),
+                        std::to_string(ExpectedCount));
+
+    for (uint32_t I = 0; I < Count; ++I) {
+      const auto *Record = (PSV.*GetMethod)(I);
+      const auto *ExpectedRecord =
+          I < ExpectedCount ? (ExpectedPSV.*GetMethod)(I) : nullptr;
+      if (!Record) {
+        EmitMismatchError(GetRecordName(Name, I), "missing record",
+                          ExpectedRecord ? GetRecordBytes(*ExpectedRecord)
+                                         : "record generated from DxilModule");
+        continue;
+      }
+      if (I < ExpectedCount && !ExpectedRecord)
+        EmitMismatchError(GetRecordName(Name, I), GetRecordBytes(*Record),
+                          "missing record generated from DxilModule");
+      VerifyTableRecord(Name, I, *Record, ExpectedRecord);
+    }
+  };
+
+  VerifyLinAlgTable(
+      "LinAlgMatrixOperationShape",
+      &DxilPipelineStateValidation::GetPSVLinAlgMatrixOperationShapeCount,
+      &DxilPipelineStateValidation::GetPSVLinAlgMatrixOperationShape,
+      VerifyRecord);
+  VerifyLinAlgTable(
+      "LinAlgMatrixConstruction",
+      &DxilPipelineStateValidation::GetPSVLinAlgMatrixConstructionCount,
+      &DxilPipelineStateValidation::GetPSVLinAlgMatrixConstruction,
+      VerifyRecordWithShapes);
+  VerifyLinAlgTable(
+      "LinAlgThreadMatrixVectorMultiply",
+      &DxilPipelineStateValidation::GetPSVLinAlgThreadMatrixVectorMultiplyCount,
+      &DxilPipelineStateValidation::GetPSVLinAlgThreadMatrixVectorMultiply,
+      VerifyRecord);
+  VerifyLinAlgTable(
+      "LinAlgWaveMatrixMultiply",
+      &DxilPipelineStateValidation::GetPSVLinAlgWaveMatrixMultiplyCount,
+      &DxilPipelineStateValidation::GetPSVLinAlgWaveMatrixMultiply,
+      VerifyRecordWithShapes);
+  VerifyLinAlgTable(
+      "LinAlgThreadGroupMatrixMultiply",
+      &DxilPipelineStateValidation::GetPSVLinAlgThreadGroupMatrixMultiplyCount,
+      &DxilPipelineStateValidation::GetPSVLinAlgThreadGroupMatrixMultiply,
+      VerifyRecordWithShapes);
+  VerifyLinAlgTable("LinAlgOuterProduct",
+                    &DxilPipelineStateValidation::GetPSVLinAlgOuterProductCount,
+                    &DxilPipelineStateValidation::GetPSVLinAlgOuterProduct,
+                    VerifyRecord);
+  VerifyLinAlgTable(
+      "LinAlgAccumulateStore",
+      &DxilPipelineStateValidation::GetPSVLinAlgAccumulateStoreCount,
+      &DxilPipelineStateValidation::GetPSVLinAlgAccumulateStore, VerifyRecord);
+}
+
 void PSVContentVerifier::Verify(unsigned ValMajor, unsigned ValMinor,
                                 unsigned PSVVersion) {
   PSVInitInfo PSVInfo(PSVVersion);
@@ -521,6 +706,8 @@ void PSVContentVerifier::Verify(unsigned ValMajor, unsigned ValMinor,
                           DM.GetEntryFunctionName());
     }
   }
+  if (PSVVersion > 3)
+    VerifyLinAlgRuntimeInfo(PSVVersion);
 
   StrTableVerifier.Verify(ValCtx);
   IndexTableVerifier.Verify(ValCtx);
@@ -607,6 +794,8 @@ bool VerifySignatureMatches(llvm::Module *pModule, DXIL::SignatureKind SigKind,
 }
 
 struct SimplePSV {
+  static bool IsDwordAligned(uint32_t Size) { return (Size & 3) == 0; }
+
   uint32_t PSVRuntimeInfoSize = 0;
   uint32_t PSVNumResources = 0;
   uint32_t PSVResourceBindInfoSize = 0;
@@ -651,7 +840,7 @@ struct SimplePSV {
       StringTableSize = GetUint32AtOffset(pPSVData, Offset);
       INCREMENT_POS(4);
       // Make sure StringTableSize is aligned to 4 bytes.
-      if ((StringTableSize & 3) != 0) {
+      if (!IsDwordAligned(StringTableSize)) {
         IsValid = false;
         return;
       }
@@ -742,7 +931,8 @@ struct SimplePSV {
         if (!Count)
           return true;
         uint32_t RecordSize = 0;
-        if (!ReadUint32(RecordSize) || RecordSize < MinimumRecordSize)
+        if (!ReadUint32(RecordSize) || !IsDwordAligned(RecordSize) ||
+            RecordSize < MinimumRecordSize)
           return false;
         if (Offset > PSVSize || Count > (PSVSize - Offset) / RecordSize)
           return false;
@@ -752,6 +942,7 @@ struct SimplePSV {
 
       uint32_t LinAlgRuntimeInfoSize = 0;
       if (!ReadUint32(LinAlgRuntimeInfoSize) ||
+          !IsDwordAligned(LinAlgRuntimeInfoSize) ||
           LinAlgRuntimeInfoSize < sizeof(PSVLinAlgRuntimeInfo0) ||
           Offset > PSVSize || LinAlgRuntimeInfoSize > PSVSize - Offset) {
         IsValid = false;

tools/clang/test/DXC/dumpPSV_LinAlg.hlsl

@@ -0,0 +1,58 @@
+// REQUIRES: dxil-1-10
+// RUN: %dxc -enable-16bit-types -E main -T cs_6_10 %s -Fo %t
+// RUN: %dxa %t -dumppsv | FileCheck %s
+
+#include <dx/linalg.h>
+using namespace dx::linalg;
+
+ByteAddressBuffer Input : register(t0);
+RWByteAddressBuffer Output : register(u0);
+RWStructuredBuffer<vector<half, 8> > VectorOutput : register(u1);
+groupshared uint8_t4_packed SharedOutput[64];
+
+using ThreadA =
+    Matrix<ComponentType::F16, 8, 4, MatrixUse::A, MatrixScope::Thread>;
+using WaveA =
+    Matrix<ComponentType::F16, 3, 4, MatrixUse::A, MatrixScope::Wave>;
+using WaveB =
+    Matrix<ComponentType::I32, 4, 5, MatrixUse::B, MatrixScope::Wave>;
+using WaveAccumulator =
+    Matrix<ComponentType::F32, 3, 5, MatrixUse::Accumulator, MatrixScope::Wave>;
+using ThreadAccumulator = Matrix<ComponentType::F32, 4, 4,
+                                 MatrixUse::Accumulator, MatrixScope::Thread>;
+
+[numthreads(4, 4, 1)]
+void main(uint Index : SV_GroupIndex) {
+  ThreadA TA =
+      ThreadA::Load<MatrixLayout::MulOptimalTranspose>(Input, 0, 0);
+  VectorOutput[Index] = Multiply<half>(TA, (vector<half, 4>)1.0h);
+
+  WaveA A = WaveA::Splat(1.0h);
+  WaveB B = WaveB::Splat(2);
+  WaveAccumulator C = Multiply<ComponentType::F32>(A, B);
+  C.Store(Output, 0, 20, MatrixLayout::RowMajor);
+  C.InterlockedAccumulate(SharedOutput, 0, 16, MatrixLayout::RowMajor);
+
+  ThreadAccumulator Outer =
+      OuterProduct<ComponentType::F32>((float4)1.0f, (float4)2.0f);
+  Outer.InterlockedAccumulate(Output, 256);
+  InterlockedAccumulate(Output, 512, (int4)Index);
+}
+
+// CHECK: LinAlgRuntimeInfoPresent: true
+// CHECK: PSVLinAlgRuntimeInfo:
+// CHECK-NEXT: MatrixOperationShapeCount: 4
+// CHECK-NEXT: MatrixConstructionCount: 3
+// CHECK-NEXT: ThreadMatrixVectorMultiplyCount: 1
+// CHECK-NEXT: WaveMatrixMultiplyCount: 1
+// CHECK-NEXT: ThreadGroupMatrixMultiplyCount: 0
+// CHECK-NEXT: OuterProductCount: 1
+// CHECK-NEXT: AccumulateStoreCount: 2
+// CHECK-NEXT: MatrixConstruction[0]: MatrixType=4, Shapes=[(0,5,4)]
+// CHECK-NEXT: MatrixConstruction[1]: MatrixType=8, Shapes=[(3,0,4)]
+// CHECK-NEXT: MatrixConstruction[2]: MatrixType=9, Shapes=[(3,5,0)]
+// CHECK-NEXT: ThreadMatrixVectorMultiply[0]: ResultType=8, MatrixType=8, VectorInputType=8, Flags=1
+// CHECK-NEXT: WaveMatrixMultiply[0]: AccumulatorType=9, MatrixAType=8, MatrixBType=4, Shapes=[(3,5,4)]
+// CHECK-NEXT: OuterProduct[0]: ResultType=9, VectorInputType=9
+// CHECK-NEXT: AccumulateStore[0]: AccumulatorType=9, Flags=3
+// CHECK-NEXT: AccumulateStore[1]: AccumulatorType=4, Flags=1

tools/clang/test/DXC/dumpPSV_LinAlgAccumulate.hlsl

@@ -0,0 +1,69 @@
+// REQUIRES: dxil-1-10
+// RUN: %dxc -enable-16bit-types -E main -T cs_6_10 %s -Fo %t
+// RUN: %dxa %t -dumppsv | FileCheck %s
+
+#include <dx/linalg.h>
+using namespace dx::linalg;
+
+RWByteAddressBuffer Output : register(u0);
+groupshared half SharedHalf[64];
+groupshared float SharedFloat[64];
+
+using ThreadHalfAccumulator = Matrix<ComponentType::F16, 2, 3,
+                                     MatrixUse::Accumulator,
+                                     MatrixScope::Thread>;
+using ThreadFloatAccumulator = Matrix<ComponentType::F32, 3, 2,
+                                      MatrixUse::Accumulator,
+                                      MatrixScope::Thread>;
+using ThreadIntAccumulator = Matrix<ComponentType::I32, 4, 4,
+                                    MatrixUse::Accumulator,
+                                    MatrixScope::Thread>;
+using WaveHalfAccumulator = Matrix<ComponentType::F16, 2, 2,
+                                   MatrixUse::Accumulator, MatrixScope::Wave>;
+using WaveFloatAccumulator = Matrix<ComponentType::F32, 2, 2,
+                                    MatrixUse::Accumulator, MatrixScope::Wave>;
+
+[numthreads(4, 4, 1)]
+void main(uint Index : SV_GroupIndex) {
+  ThreadHalfAccumulator HalfOuter =
+      OuterProduct<ComponentType::F16>((vector<half, 2>)1.0h,
+                                       (vector<half, 3>)2.0h);
+  HalfOuter.InterlockedAccumulate(Output, 0);
+
+  ThreadFloatAccumulator FloatOuter =
+      OuterProduct<ComponentType::F32>((vector<half, 3>)3.0h,
+                                       (vector<half, 2>)4.0h);
+  FloatOuter.InterlockedAccumulate(Output, 64);
+
+  ThreadIntAccumulator IntOuter =
+      OuterProduct<ComponentType::I32>((int4)5, (int4)6);
+  IntOuter.InterlockedAccumulate(Output, 128);
+
+  WaveHalfAccumulator WaveHalf = WaveHalfAccumulator::Splat(7.0h);
+  WaveHalf.InterlockedAccumulate(Output, 192, 4, MatrixLayout::RowMajor);
+  WaveHalf.InterlockedAccumulate(SharedHalf, 0, 8, MatrixLayout::RowMajor);
+
+  WaveFloatAccumulator WaveFloat = WaveFloatAccumulator::Splat(8.0f);
+  WaveFloat.InterlockedAccumulate(SharedFloat, 0, 4, MatrixLayout::RowMajor);
+
+  InterlockedAccumulate(Output, 256, (vector<int64_t, 2>)Index);
+}
+
+// CHECK: LinAlgRuntimeInfoPresent: true
+// CHECK: PSVLinAlgRuntimeInfo:
+// CHECK-NEXT: MatrixOperationShapeCount: 1
+// CHECK-NEXT: MatrixConstructionCount: 2
+// CHECK-NEXT: ThreadMatrixVectorMultiplyCount: 0
+// CHECK-NEXT: WaveMatrixMultiplyCount: 0
+// CHECK-NEXT: ThreadGroupMatrixMultiplyCount: 0
+// CHECK-NEXT: OuterProductCount: 3
+// CHECK-NEXT: AccumulateStoreCount: 4
+// CHECK-NEXT: MatrixConstruction[0]: MatrixType=8, Shapes=[(2,2,0)]
+// CHECK-NEXT: MatrixConstruction[1]: MatrixType=9, Shapes=[(2,2,0)]
+// CHECK-NEXT: OuterProduct[0]: ResultType=8, VectorInputType=8
+// CHECK-NEXT: OuterProduct[1]: ResultType=9, VectorInputType=8
+// CHECK-NEXT: OuterProduct[2]: ResultType=4, VectorInputType=4
+// CHECK-NEXT: AccumulateStore[0]: AccumulatorType=8, Flags=3
+// CHECK-NEXT: AccumulateStore[1]: AccumulatorType=9, Flags=3
+// CHECK-NEXT: AccumulateStore[2]: AccumulatorType=4, Flags=1
+// CHECK-NEXT: AccumulateStore[3]: AccumulatorType=6, Flags=1