Add verifier for migraphx.reshape (#1999)

* Add verifier for MIGraphX Reshape op

* Update for values of -1

* Clang-format

* Attend to review comments

* Make more changes based on review comments

* Fix outDim handling

Author: justinrosner

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

@@ -439,18 +439,25 @@ def MIGraphX_TransposeOp :
 }
 
 def MIGraphX_ReshapeOp :
-    MIGraphX_Op<"reshape">,
+    MIGraphX_Op<"reshape", [AllElementTypesMatch<["input", "output"]>]>,
     Arguments<(ins AnyMIXRShaped:$input,
                    I64ArrayAttr:$dims
                    )>,
 	Results<(outs AnyMIXRShaped:$output)> {
   let summary = "reshape a tensor";
   let description = [{
     The `migraphx.reshape` op.
+
+    `dims` can contain values of -1 and 0. A value of -1 means to infer this 
+    dimension from the others. i.e., If we have an output shape of
+    `<4x2xf32, 0x1>` and the dims are `[4, -1]`, then the second dimension will
+    be inferred to be 2. A value of 0 means explicit zero-length dimension.
   }];
   let assemblyFormat = [{
     $input attr-dict `:` type($input) `->` type($output)
   }];
+
+  let hasVerifier = 1;
 }
 
 def MIGraphX_SliceOp :

mlir/lib/Dialect/MIGraphX/IR/MIGraphX.cpp

@@ -313,6 +313,61 @@ LogicalResult LiteralOp::verify() {
   return success();
 }
 
+LogicalResult ReshapeOp::verify() {
+  MIXRShapedType inputType = getInput().getType();
+  MIXRShapedType outType = getOutput().getType();
+  ArrayAttr dimsAttr = getDims();
+
+  // Dynamic shapes are not currently supported
+  if (!inputType.hasStaticShape())
+    return emitOpError("Dynamic shapes are not supported");
+
+  if (dimsAttr.size() != outType.getRank())
+    return emitOpError("number of dims (")
+           << dimsAttr.size() << ") does not match result rank ("
+           << outType.getRank() << ")";
+
+  // Check that there is only a single -1 value
+  int missingDims = llvm::count_if(
+      dimsAttr.getAsRange<IntegerAttr>(),
+      [](IntegerAttr a) { return a.getInt() == -1; });
+  if (missingDims > 1)
+    return emitOpError("expected at most one target dimension to be -1");
+
+  // Check how many zero dimensions there are
+  int numZeros = llvm::count_if(
+      dimsAttr.getAsRange<IntegerAttr>(),
+      [](IntegerAttr a) { return a.getInt() == 0; });
+
+  if (missingDims > 0 && numZeros > 0)
+    return emitOpError("Cannot mix missing dimensions with zero dimension");
+
+  // Compare dimension values to output shape
+  for (auto [dimVal, outDim] : llvm::zip(dimsAttr, outType.getShape())) {
+    int64_t dimValue = cast<IntegerAttr>(dimVal).getInt();
+    // We cannot handle negative dims values that aren't -1 
+    if (dimValue < -1 ) {
+      return emitOpError("Non -1 negative values are not supported");
+    }
+
+    // Output dimensions can't be negative
+    if (outDim < 0)
+      return emitOpError("Negative output dimensions are not supported");
+
+    // Per-dimension consistency
+    if (dimValue >= 0 && outDim != dimValue)
+      return emitOpError("dimValue: ")
+             << dimValue << " inconsistent with result dimension " << outDim;
+  }
+
+  // Check that the number of elements in the input and output types match
+  int64_t inputElements = inputType.getNumElements();
+  if (inputElements != outType.getNumElements())
+    return emitOpError("input and output element counts do not match");
+
+  return success();
+}
+
 LogicalResult UnpackOp::verify() {
   MIXRShapedType inType = getIn().getType();
   MIXRShapedType outType = getOut().getType();

mlir/test/Dialect/MIGraphX/invalid.mlir

@@ -0,0 +1,44 @@
+// RUN: rocmlir-opt %s -split-input-file -verify-diagnostics
+
+func.func @mlir_reshape_inconsistent_dims(%arg0: !migraphx.shaped<4096x4096xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op dimValue: 64 inconsistent with result dimension 4096}}
+  %0 = migraphx.reshape %arg0 {dims = [64, 128]} : <4096x4096xf16, 0x1> -> <4096x4096xf16, 16536x2>
+  return
+}
+
+func.func @mlir_reshape_dynamic_shape(%arg0: !migraphx.shaped<4096x?xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op Dynamic shapes are not supported}}
+  %0 = migraphx.reshape %arg0 {dims = [4096, 4096]} : <4096x?xf16, 0x1> -> <4096x?xf16, 16536x2>
+  return
+}
+
+func.func @mlir_reshape_rank(%arg0: !migraphx.shaped<4096x4096xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op number of dims (3) does not match result rank (2)}}
+  %0 = migraphx.reshape %arg0 {dims = [1, 4096, 4096]} : <4096x4096xf16, 0x1> -> <4096x4096xf16, 16536x2>
+  return
+}
+
+func.func @mlir_num_input_elements(%arg0: !migraphx.shaped<2x4xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op input and output element counts do not match}}
+  %0 = migraphx.reshape %arg0 {dims = [3, 5]} : <2x4xf16, 0x1> -> <3x5xf16, 0x1>
+  return
+}
+
+func.func @mlir_element_type(%arg0: !migraphx.shaped<2x4xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op failed to verify that all of {input, output} have same element type}}
+  %0 = migraphx.reshape %arg0 {dims = [4, 2]} : <2x4xf16, 0x1> -> <4x2xf32, 0x1>
+  return
+}
+
+func.func @mlir_multiple_neg_one(%arg0: !migraphx.shaped<2x4xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op expected at most one target dimension to be -1}}
+  %0 = migraphx.reshape %arg0 {dims = [-1, -1]} : <2x4xf16, 0x1> -> <4x2xf16, 0x1>
+  return
+}
+
+func.func @mlir_neg_one_with_zero(%arg0: !migraphx.shaped<2x4xf16, 0x1>) {
+  // expected-error@+1 {{'migraphx.reshape' op Cannot mix missing dimensions with zero dimension}}
+  %0 = migraphx.reshape %arg0 {dims = [0, -1]} : <2x4xf16, 0x1> -> <4x2xf16, 0x1>
+  return
+}
+

mlir/test/fusion/mixr-attention-realtype-int4-dequantizelinear.mlir

@@ -15,8 +15,8 @@ module {
     %0 = migraphx.unpack %arg3 {axis = 1 : i64} : <4096x2048xui8, 2048x1> -> <4096x4096xi8, 4096x1>
     %1 = migraphx.broadcast %arg1 {axis = 0 : i64, out_lens = [4096, 4096]} : <4096xf16, 1> -> <4096x4096xf16, 0x1>
     %2 = migraphx.broadcast %arg2 {axis = 0 : i64, out_lens = [4096, 4096]} : <4096xf16, 1> -> <4096x4096xf16, 0x1>
-    %3 = migraphx.reshape %1 {dims = [64, 128]} : <4096x4096xf16, 0x1> -> <4096x4096xf16, 16536x2>
-    %4 = migraphx.reshape %2 {dims = [64, 128]} : <4096x4096xf16, 0x1> -> <4096x4096xf16, 16536x2>
+    %3 = migraphx.reshape %1 {dims = [4096, 4096]} : <4096x4096xf16, 0x1> -> <4096x4096xf16, 16536x2>
+    %4 = migraphx.reshape %2 {dims = [4096, 4096]} : <4096x4096xf16, 0x1> -> <4096x4096xf16, 16536x2>
     %5 = migraphx.dequantizelinear %0, %3, %4 : <4096x4096xi8, 4096x1>, <4096x4096xf16, 16536x2>, !migraphx.shaped<4096x4096xf16, 16536x2> -> <4096x4096xf16, 4096x1>
     %6 = migraphx.dot %2, %5 : <4096x4096xf16, 0x1>, <4096x4096xf16, 4096x1> -> <4096x4096xf16, 4096x1>
     %7 = migraphx.softmax %6 {axis = 1 : i64} : <4096x4096xf16, 4096x1> -> <4096x4096xf16, 4096x1>