[PATCH] D147438: [Matrix] Add shape verification.

[Florian Hahn via Phabricator via llvm-commits]

fhahn created this revision.
fhahn added reviewers: anemet, thegameg.
Herald added subscribers: StephenFan, tschuett, hiraditya.
Herald added a project: All.
fhahn requested review of this revision.
Herald added a project: LLVM.

At the moment, lower-matrix-intrinsics accepts mis-matches between
shapes for operations. See shape-verification.ll for an example where
@llvm.matrix.column.major.load specifies 6x1 and then the use
(@llvm.matrix.multiply) specifies the operand to have 1x6.

This patch adds verification for shapes to check if shapes match.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D147438

Files:
  llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
  llvm/test/Transforms/LowerMatrixIntrinsics/shape-verification.ll


Index: llvm/test/Transforms/LowerMatrixIntrinsics/shape-verification.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/LowerMatrixIntrinsics/shape-verification.ll
@@ -0,0 +1,16 @@
+; RUN: not --crash opt -passes='lower-matrix-intrinsics' -verify-matrix-shapes=true -S %s 2>&1 | FileCheck --check-prefix=VERIFY %s
+; RUN: opt -passes='lower-matrix-intrinsics' -verify-matrix-shapes=false -S %s 2>&1 | FileCheck --check-prefix=NOVERIFY %s
+
+; VERIFY: Conflicting shapes (6x1 vs 1x6)
+; NOVERIFY-NOT: Conflicting shapes
+
+define <1 x float> @intrinsic_column_major_load_dot_product_float_v6(ptr %lhs_address, ptr %rhs_address) {
+entry:
+  %lhs = tail call fast <6 x float> @llvm.matrix.column.major.load.v6f32.i64(ptr nonnull align 4 %lhs_address, i64 6, i1 false, i32 6, i32 1)
+  %rhs = tail call fast <6 x float> @llvm.matrix.column.major.load.v6f32.i64(ptr nonnull align 4 %rhs_address, i64 1, i1 false, i32 1, i32 6)
+  %result = tail call fast <1 x float> @llvm.matrix.multiply.v1f32.v6f32.v6f32(<6 x float> %lhs, <6 x float> %rhs, i32 1, i32 6, i32 1)
+  ret <1 x float> %result
+}
+
+declare <6 x float> @llvm.matrix.column.major.load.v6f32.i64(ptr nonnull align 4, i64, i1, i32, i32)
+declare <1 x float> @llvm.matrix.multiply.v1f32.v6f32.v6f32(<6 x float>, <6 x float>, i32, i32, i32)
Index: llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
+++ llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
@@ -72,6 +72,16 @@
     cl::desc("Allow the use of FMAs if available and profitable. This may "
              "result in different results, due to less rounding error."));
 
+static cl::opt<bool>
+    VerifyShapeInfo("verify-matrix-shapes", cl::Hidden,
+                    cl::desc("Enable/disable matrix shape verification."),
+#ifdef NDEBUG
+                    cl::init(false)
+#else
+                    cl::init(true)
+#endif
+    );
+
 enum class MatrixLayoutTy { ColumnMajor, RowMajor };
 
 static cl::opt<MatrixLayoutTy> MatrixLayout(
@@ -535,6 +545,15 @@
 
     auto SIter = ShapeMap.find(V);
     if (SIter != ShapeMap.end()) {
+      if (VerifyShapeInfo && (SIter->second.NumRows != Shape.NumRows ||
+                              SIter->second.NumColumns != Shape.NumColumns)) {
+        errs() << "Conflicting shapes (" << SIter->second.NumRows << "x"
+               << SIter->second.NumColumns << " vs " << Shape.NumRows << "x"
+               << Shape.NumColumns << ") for " << *V << "\n";
+        report_fatal_error(
+            "Matrix shape verification failed, compilation aborted!");
+      }
+
       LLVM_DEBUG(dbgs() << "  not overriding existing shape: "
                         << SIter->second.NumRows << " "
                         << SIter->second.NumColumns << " for " << *V << "\n");


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D147438.510457.patch
Type: text/x-patch
Size: 2917 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230403/aec8fcbb/attachment.bin>