fix: restrict raising dot_general to syrk to 2D tensors (#1666)

Author: avik-pal

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -25884,6 +25884,14 @@ struct DotGeneralToSyrk
     auto lhs = op.getLhs();
     auto rhs = op.getRhs();
 
+    auto lhsType = cast<RankedTensorType>(lhs.getType());
+    auto rhsType = cast<RankedTensorType>(rhs.getType());
+    auto outType = cast<RankedTensorType>(op.getResult().getType());
+    if (lhsType.getRank() != 2 || rhsType.getRank() != 2 ||
+        outType.getRank() != 2) {
+      return failure();
+    }
+
     if (dotDims.getLhsBatchingDimensions().size() != 0 ||
         dotDims.getRhsBatchingDimensions().size() != 0) {
       return failure();

test/lit_tests/dotgeneral_to_syrk.mlir

@@ -71,3 +71,15 @@ func.func @main4(%arg0: tensor<64x32xf32>) -> tensor<64x64xf32> {
 // CHECK-NEXT:   %0 = enzymexla.blas.syrk %arg0, %cst_1, %cst_0, %cst {fill, uplo = #enzymexla.uplo<F>} : (tensor<64x32xf32>, tensor<64x64xf32>, tensor<f32>, tensor<f32>) -> tensor<64x64xf32>
 // CHECK-NEXT:   return %0 : tensor<64x64xf32>
 // CHECK-NEXT: }
+
+func.func @fail1(%arg0: tensor<5x2xf32>) -> tensor<f32> {
+  %0 = stablehlo.reshape %arg0 : (tensor<5x2xf32>) -> tensor<10xf32>
+  %1 = stablehlo.dot_general %0, %0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<10xf32>, tensor<10xf32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK: func.func @fail1(%arg0: tensor<5x2xf32>) -> tensor<f32> {
+// CHECK-NEXT:   %0 = stablehlo.reshape %arg0 : (tensor<5x2xf32>) -> tensor<10xf32>
+// CHECK-NEXT:   %1 = stablehlo.dot_general %0, %0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<10xf32>, tensor<10xf32>) -> tensor<f32>
+// CHECK-NEXT:   return %1 : tensor<f32>
+// CHECK-NEXT: }