Merged PR 44: Fix CreateTransposeShuffle in HLMatrixLowerPass

Fix CreateTransposeShuffle in HLMatrixLowerPass

lib/HLSL/HLMatrixLowerPass.cpp

@@ -1265,12 +1265,12 @@ void HLMatrixLowerPass::TrivialMatReplace(Value *matVal,
     }
 }
 
-static Instruction *CreateTransposeShuffle(IRBuilder<> &Builder, Value *vecVal, unsigned row, unsigned col) {
-  SmallVector<int, 16> castMask(col * row);
+static Instruction *CreateTransposeShuffle(IRBuilder<> &Builder, Value *vecVal, unsigned toRows, unsigned toCols) {
+  SmallVector<int, 16> castMask(toCols * toRows);
   unsigned idx = 0;
-  for (unsigned c = 0; c < col; c++)
-    for (unsigned r = 0; r < row; r++)
-      castMask[idx++] = r * col + c;
+  for (unsigned r = 0; r < toRows; r++)
+    for (unsigned c = 0; c < toCols; c++)
+      castMask[idx++] = c * toRows + r;
   return cast<Instruction>(
     Builder.CreateShuffleVector(vecVal, vecVal, castMask));
 }

lib/HLSL/HLOperationLower.cpp

@@ -6745,12 +6745,12 @@ static Instruction *BitCastValueOrPtr(Value* V, Instruction *Insert, Type *Ty, b
   }
 }
 
-static Instruction *CreateTransposeShuffle(IRBuilder<> &Builder, Value *vecVal, unsigned row, unsigned col) {
-  SmallVector<int, 16> castMask(col * row);
+static Instruction *CreateTransposeShuffle(IRBuilder<> &Builder, Value *vecVal, unsigned toRows, unsigned toCols) {
+  SmallVector<int, 16> castMask(toCols * toRows);
   unsigned idx = 0;
-  for (unsigned c = 0; c < col; c++)
-    for (unsigned r = 0; r < row; r++)
-      castMask[idx++] = r * col + c;
+  for (unsigned r = 0; r < toRows; r++)
+    for (unsigned c = 0; c < toCols; c++)
+      castMask[idx++] = c * toRows + r;
   return cast<Instruction>(
     Builder.CreateShuffleVector(vecVal, vecVal, castMask));
 }
@@ -6834,11 +6834,11 @@ void TranslateHLBuiltinOperation(Function *F, HLOperationLowerHelper &helper,
           IRBuilder<> Builder(CI);
           HLCastOpcode opcode = static_cast<HLCastOpcode>(hlsl::GetHLOpcode(CI));
           bool bTranspose = false;
-          bool bColSource = false;
+          bool bColDest = false;
           switch (opcode) {
-          case HLCastOpcode::ColMatrixToRowMatrix:
-            bColSource = true;
           case HLCastOpcode::RowMatrixToColMatrix:
+            bColDest = true;
+          case HLCastOpcode::ColMatrixToRowMatrix:
             bTranspose = true;
           case HLCastOpcode::ColMatrixToVecCast:
           case HLCastOpcode::RowMatrixToVecCast: {
@@ -6849,7 +6849,7 @@ void TranslateHLBuiltinOperation(Function *F, HLOperationLowerHelper &helper,
             if (bTranspose) {
               unsigned row, col;
               HLMatrixLower::GetMatrixInfo(matVal->getType(), col, row);
-              if (bColSource) std::swap(row, col);
+              if (bColDest) std::swap(row, col);
               vecVal = CreateTransposeShuffle(Builder, vecVal, row, col);
             }
             CI->replaceAllUsesWith(vecVal);

tools/clang/test/CodeGenHLSL/quick-test/mat_col_default1.hlsl

@@ -0,0 +1,20 @@
+// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+
+// Make sure get cb0[1].y and cb0[1].z.
+// CHECK: call %dx.types.CBufRet.f32 @dx.op.cbufferLoadLegacy.f32
+// CHECK: extractvalue %dx.types.CBufRet.f32 {{.*}}, 1
+// CHECK: extractvalue %dx.types.CBufRet.f32 {{.*}}, 2
+
+cbuffer Transform : register(b0)
+{
+  float4 transformRows[3];
+}
+
+float2 main(int i : A) : SV_TARGET
+{
+  float3x4 mat;
+  mat[0] = transformRows[0];
+  mat[1] = transformRows[1];
+  mat[2] = transformRows[2];
+  return mat[1].yz;
+}