[Verifier] Verify matrix dimensions operands match vector size.

[fhahn]

Summary

This patch adds checks to the verifier to ensure the dimension arguments
passed to the matrix intrinsics match the vector types for their
arugments/return values.

[anemet]

Nice, LGTM!

[SjoerdMeijer]

llvm/lib/IR/Verifier.cpp

4805

Quick query on this and the semantics:

declare vectorty @llvm.matrix.multiply.*(vectorty %A, vectorty %B, i32 <OuterRows>, i32 <Inner>, i32 <OuterColumns>)

do we expect the element types of vectors %A and %B to be same, and do we need to check this?

[fhahn]

llvm/lib/IR/Verifier.cpp

4805

Yes, the element types of all types must match currently, but I think it is neither checked in the verifier nor explicit in the LangRef.

To generate code for llvm.aarch64.neon.udot & co, there probably needs to be a way to have different element type widths for result and source operands.

[SjoerdMeijer]

llvm/lib/IR/Verifier.cpp

4805

Yes, the element types of all types must match currently, but I think it is neither checked in the verifier nor explicit in the LangRef.

I started looking at the matrix support, getting up to speed with it, and this is where I started and the first thing I noticed. Was just asking about that here as a sanity check. I wouldn't mind putting up a patch for that if that's helpful. Probably the least we can do for not is to check if we are not mixing integers and float types, and then we also need to add that to LangRef and be explicit about that.

[fhahn]

llvm/lib/IR/Verifier.cpp

4805

I wouldn't mind putting up a patch for that if that's helpful.

That would be great. I think things will fall apart/miscompile if the element types differ at the moment.

[SjoerdMeijer]

llvm/lib/IR/Verifier.cpp

4805

cool, will do.

FYI: I started with committing a NFC patch adding some more negative tests: rG0fc17e9edc8f

llvm/lib/IR/Verifier.cpp

case Intrinsic::llrint: {
     break;
   }
   case Intrinsic::bswap: {
     Type *Ty = Call.getType();
     unsigned Size = Ty->getScalarSizeInBits();
     Assert(Size % 16 == 0, "bswap must be an even number of bytes", &Call);
     break;
   }
+  case Intrinsic::matrix_multiply:
+  case Intrinsic::matrix_transpose:
+  case Intrinsic::matrix_columnwise_load:
+  case Intrinsic::matrix_columnwise_store: {
+    ConstantInt *NumRows;
+    ConstantInt *NumColumns;
+    VectorType *TypeToCheck;
+    switch (ID) {
+    case Intrinsic::matrix_multiply:
+      NumRows = cast<ConstantInt>(Call.getArgOperand(2));
+      NumColumns = cast<ConstantInt>(Call.getArgOperand(4));
+      TypeToCheck = cast<VectorType>(Call.getType());
+      break;
+    case Intrinsic::matrix_transpose:
+      NumRows = cast<ConstantInt>(Call.getArgOperand(1));
+      NumColumns = cast<ConstantInt>(Call.getArgOperand(2));
+      TypeToCheck = cast<VectorType>(Call.getType());
+      break;
+    case Intrinsic::matrix_columnwise_load:
+      NumRows = cast<ConstantInt>(Call.getArgOperand(2));
+      NumColumns = cast<ConstantInt>(Call.getArgOperand(3));
+      TypeToCheck = cast<VectorType>(Call.getType());
+      break;
+    case Intrinsic::matrix_columnwise_store:
+      NumRows = cast<ConstantInt>(Call.getArgOperand(3));
+      NumColumns = cast<ConstantInt>(Call.getArgOperand(4));
+      TypeToCheck = cast<VectorType>(Call.getArgOperand(0)->getType());
+      break;
+    default:
+      llvm_unreachable("unexpected intrinsic");
+    }
+    Assert(TypeToCheck->getNumElements() ==
+               NumRows->getZExtValue() * NumColumns->getZExtValue(),
+           "result of a matrix operation does not fit in the returned vector");
+    break;
+  }
   };
 }
 
 /// Carefully grab the subprogram from a local scope.
 ///
 /// This carefully grabs the subprogram from a local scope, avoiding the
 /// built-in assertions that would typically fire.
 static DISubprogram *getSubprogram(Metadata *LocalScope) {

llvm/test/Verifier/matrix-intrinsics.ll

+; RUN: not llvm-as < %s -o /dev/null 2>&1 | FileCheck %s
+
+declare <4 x float> @llvm.matrix.transpose.v4f32(<4 x float>, i32, i32)
+define <4 x float> @transpose(<4 x float> %m) {
+; CHECK: assembly parsed, but does not verify as correct!
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+  %result.1 = call <4 x float> @llvm.matrix.transpose.v4f32(<4 x float> %m, i32 3, i32 2)
+  %result.2 = call <4 x float> @llvm.matrix.transpose.v4f32(<4 x float> %result.1, i32 2, i32 1)
+  ret <4 x float> %result.2
+}
+
+declare <4 x float> @llvm.matrix.multiply.v4f32.v4f32.v4f32(<4 x float>, <4 x float>, i32, i32, i32)
+define <4 x float> @multiply(<4 x float> %m) {
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+  %result.1 = call <4 x float> @llvm.matrix.multiply.v4f32.v4f32.v4f32(<4 x float> %m, <4 x float> %m, i32 3, i32 2, i32 2)
+  %result.2 = call <4 x float> @llvm.matrix.multiply.v4f32.v4f32.v4f32(<4 x float> %result.1, <4 x float> %m, i32 2, i32 2, i32 1)
+  ret <4 x float> %result.2
+}
+
+declare <4 x float> @llvm.matrix.columnwise.load.v4f32.p0v4f32(<4 x float>*, i32, i32, i32)
+declare <6 x float> @llvm.matrix.columnwise.load.v6f32.p0v6f32(<6 x float>*, i32, i32, i32)
+define <4 x float> @columnwise_load(<4 x float>* %m, <6 x float>* %n) {
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+  %result.1 = call <4 x float> @llvm.matrix.columnwise.load.v4f32.p0v4f32(<4 x float>* %m, i32 2, i32 1, i32 2)
+  %result.2 = call <6 x float> @llvm.matrix.columnwise.load.v6f32.p0v6f32(<6 x float>* %n, i32 2, i32 3, i32 3)
+  ret <4 x float> %result.1
+}
+
+declare void @llvm.matrix.columnwise.store.v4f32.p0v4f32(<4 x float>, <4 x float>*, i32, i32, i32)
+declare void @llvm.matrix.columnwise.store.v6f32.p0v6f32(<6 x float>, <6 x float>*, i32, i32, i32)
+define void @columnwise_store(<4 x float>* %m, <6 x float>* %n) {
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+; CHECK-NEXT: result of a matrix operation does not fit in the returned vector
+  call void @llvm.matrix.columnwise.store.v4f32.p0v4f32(<4 x float> zeroinitializer, <4 x float>* %m, i32 2, i32 1, i32 2)
+  call void @llvm.matrix.columnwise.store.v6f32.p0v6f32(<6 x float> zeroinitializer, <6 x float>* %n, i32 2, i32 3, i32 3)
+  ret void
+}