[Matrix] Fix miscompile for NT matmul if the transpose has other use

[anemet]

Summary

We should only add the fake lowering entry for the matrix remark if the
transpose is not lowered on its own. MapVector::insert is used to insert
the entry during proper lowering which does not overwrite the fake entry in
the map.

We actually had test coverage for this but the reference output code was
wrong; it was storing undef rather than the transposed column.

[fhahn]

LGTM, thanks!

[anemet]

In the committed version, I added a new assert that should have caught this.

[MaskRay]

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

1144

Hi, there was a -Wunused-variable in -DLLVM_ENABLE_ASSERTIONS=off build

Fixed by 3b181568db8efc30b79466bba4aa334f86a49877

[anemet]

thanks @MaskRay

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

public:
 
   /// Cache \p Matrix as result of \p Inst and update the uses of \p Inst. For
   /// users with shape information, there's nothing to do: they will use the
   /// cached value when they are lowered. For other users, \p Matrix is
   /// flattened and the uses are updated to use it. Also marks \p Inst for
   /// deletion.
   void finalizeLowering(Instruction *Inst, MatrixTy Matrix,
                         IRBuilder<> &Builder) {
-    Inst2ColumnMatrix.insert(std::make_pair(Inst, Matrix));
+    auto inserted = Inst2ColumnMatrix.insert(std::make_pair(Inst, Matrix));
+    assert(inserted.second && "multiple matrix lowering mapping");
 
     ToRemove.push_back(Inst);
     Value *Flattened = nullptr;
     for (Use &U : llvm::make_early_inc_range(Inst->uses())) {
       if (ShapeMap.find(U.getUser()) == ShapeMap.end()) {
         if (!Flattened)
           Flattened = Matrix.embedInVector(Builder);
         U.set(Flattened);
if (MatrixLayout == MatrixLayoutTy::ColumnMajor
 
       emitMatrixMultiply(Result, MA, MB, Builder, false, true,
                          getFastMathFlags(MatMul));
 
       FusedInsts.insert(MatMul);
       if (Transpose->hasOneUse()) {
         FusedInsts.insert(cast<Instruction>(Transpose));
         ToRemove.push_back(cast<Instruction>(Transpose));
-      }
-      finalizeLowering(MatMul, Result, Builder);
         // TODO: add a fake entry for the folded instruction so that this is
         // included in the expression in the remark.
         Inst2ColumnMatrix[Transpose] = MatrixTy(M, C, EltType);
+      }
+      finalizeLowering(MatMul, Result, Builder);
       return;
     }
 
     if (!MatMul->hasOneUse() || MatrixLayout != MatrixLayoutTy::ColumnMajor)
       return;
 
     // Lower {ld, ld} -> matmul -> st chains.  No need to call finalizeLowering
     // since the single store user will be lowered as part of this.

llvm/test/Transforms/LowerMatrixIntrinsics/multiply-right-transpose.ll

 ; CHECK-NEXT:    [[TMP6:%.*]] = extractelement <2 x double> [[SPLIT39]], i64 1
 ; CHECK-NEXT:    [[TMP7:%.*]] = insertelement <3 x double> undef, double [[TMP6]], i64 0
 ; CHECK-NEXT:    [[TMP8:%.*]] = extractelement <2 x double> [[SPLIT40]], i64 1
 ; CHECK-NEXT:    [[TMP9:%.*]] = insertelement <3 x double> [[TMP7]], double [[TMP8]], i64 1
 ; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <2 x double> [[SPLIT41]], i64 1
 ; CHECK-NEXT:    [[TMP11:%.*]] = insertelement <3 x double> [[TMP9]], double [[TMP10]], i64 2
 ; CHECK-NEXT:    [[TMP12:%.*]] = bitcast <6 x double>* [[P:%.*]] to double*
 ; CHECK-NEXT:    [[VEC_CAST:%.*]] = bitcast double* [[TMP12]] to <3 x double>*
-; CHECK-NEXT:    store <3 x double> undef, <3 x double>* [[VEC_CAST]], align 16
+; CHECK-NEXT:    store <3 x double> [[TMP5]], <3 x double>* [[VEC_CAST]], align 16
 ; CHECK-NEXT:    [[VEC_GEP:%.*]] = getelementptr double, double* [[TMP12]], i64 3
 ; CHECK-NEXT:    [[VEC_CAST42:%.*]] = bitcast double* [[VEC_GEP]] to <3 x double>*
-; CHECK-NEXT:    store <3 x double> undef, <3 x double>* [[VEC_CAST42]], align 8
+; CHECK-NEXT:    store <3 x double> [[TMP11]], <3 x double>* [[VEC_CAST42]], align 8
 ; CHECK-NEXT:    [[SPLIT:%.*]] = shufflevector <6 x double> [[A:%.*]], <6 x double> poison, <2 x i32> <i32 0, i32 1>
 ; CHECK-NEXT:    [[SPLIT1:%.*]] = shufflevector <6 x double> [[A]], <6 x double> poison, <2 x i32> <i32 2, i32 3>
 ; CHECK-NEXT:    [[SPLIT2:%.*]] = shufflevector <6 x double> [[A]], <6 x double> poison, <2 x i32> <i32 4, i32 5>
 ; CHECK-NEXT:    [[SPLIT3:%.*]] = shufflevector <6 x double> [[B]], <6 x double> poison, <2 x i32> <i32 0, i32 1>
 ; CHECK-NEXT:    [[SPLIT4:%.*]] = shufflevector <6 x double> [[B]], <6 x double> poison, <2 x i32> <i32 2, i32 3>
 ; CHECK-NEXT:    [[SPLIT5:%.*]] = shufflevector <6 x double> [[B]], <6 x double> poison, <2 x i32> <i32 4, i32 5>
 ; CHECK-NEXT:    [[BLOCK:%.*]] = shufflevector <2 x double> [[SPLIT]], <2 x double> poison, <1 x i32> zeroinitializer
 ; CHECK-NEXT:    [[TMP13:%.*]] = extractelement <2 x double> [[SPLIT3]], i64 0