[NFC] Add elementwise trait on MIGraphX dialect ops (#2339)

* Add Elementwise trait to MIGraphX dialect elementwise ops

Mark all MIGraphX elementwise operations (unary, binary, clip, where,
convert) with the MLIR Elementwise trait. This enables programmatic
querying of elementwise ops via op.hasTrait<OpTrait::Elementwise>()
instead of exhaustive isa<> lists.

Also add AllShapesMatch<["inA", "output"]> to unary ops and convert to
enforce that elementwise unary operations preserve shape.

Author: umangyadav

mlir/include/mlir/Conversion/LinalgToRock/LinalgToRock.h

@@ -28,7 +28,7 @@ void populateLinalgToRockConversionPattern(RewritePatternSet &pattern,
 
 /// A tensor.insert_slice is said to be a rock.expand_strides
 bool isRockExpandStride(tensor::InsertSliceOp op);
-}
+} // namespace rock
 } // namespace mlir
 
 #endif

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

@@ -58,7 +58,8 @@ def MIGraphX_LiteralOp : MIGraphX_Op<"literal",
 class MIGraphX_ElementwiseBinaryOp<string name, list<Trait> traits = []>
     : MIGraphX_Op<
           name, !listconcat(
-                    traits, [AllElementTypesMatch<["inA", "inB", "output"]>])>,
+                    traits, [Elementwise,
+                             AllElementTypesMatch<["inA", "inB", "output"]>])>,
       Arguments<(ins AnyMIXRShaped:$inA, AnyMIXRShaped:$inB)>,
       Results<(outs AnyMIXRShaped:$output)> {
   let summary = "Elementwise " # name # " of two shaped values with broadcast";
@@ -94,12 +95,14 @@ def MIGraphX_Equal : MIGraphX_ElementwiseBinaryOp<"equal"> {
   }];
 }
 
-def MIGraphX_ClipOp :
-    MIGraphX_Op<"clip">,
-    Arguments<(ins AnyMIXRShaped:$x,
-      AnyMIXRShaped:$minVals,
-      AnyMIXRShaped:$maxVals)>,
-    Results<(outs AnyMIXRShaped:$output)> {
+def MIGraphX_ClipOp
+    : MIGraphX_Op<
+          "clip", [Elementwise,
+                   AllElementTypesMatch<["x", "minVals", "maxVals", "output"]>,
+                   AllShapesMatch<["x", "minVals", "maxVals", "output"]>]>,
+      Arguments<(ins AnyMIXRShaped:$x, AnyMIXRShaped:$minVals,
+          AnyMIXRShaped:$maxVals)>,
+      Results<(outs AnyMIXRShaped:$output)> {
   let summary = "Elementwise clip";
   let description = [{
     Elementwise clip: output = min(max(x, minVals), maxVals)
@@ -113,7 +116,8 @@ def MIGraphX_ClipOp :
 // Note: when lowering to kernel calls, MIGraphX represents booleans as i8.
 // Keep that logic here.
 def MIGraphX_WhereOp
-    : MIGraphX_Op<"where", [AllElementTypesMatch<["inA", "inB", "output"]>,
+    : MIGraphX_Op<"where", [Elementwise,
+                            AllElementTypesMatch<["inA", "inB", "output"]>,
                             AllShapesMatch<["inA", "inB", "output", "cond"]>]>,
       Arguments<(ins MIXRShapedOf<[I8, SI8, UI8]>:$cond, AnyMIXRShaped:$inA,
           AnyMIXRShaped:$inB)>,
@@ -130,10 +134,10 @@ def MIGraphX_WhereOp
 
 // Elementwise unary operations
 
-def MIGraphX_ConvertOp :
-    MIGraphX_Op<"convert">,
-    Arguments<(ins AnyMIXRShaped:$inA)>,
-	  Results<(outs AnyMIXRShaped:$output)> {
+def MIGraphX_ConvertOp
+    : MIGraphX_Op<"convert", [Elementwise, AllShapesMatch<["inA", "output"]>]>,
+      Arguments<(ins AnyMIXRShaped:$inA)>,
+      Results<(outs AnyMIXRShaped:$output)> {
   let summary = "Elementwise type conversion";
   let description = [{
     Type conversion. Due to impedance mismatches between MIGraphX and Tosa,
@@ -142,10 +146,12 @@ def MIGraphX_ConvertOp :
   let assemblyFormat = "$inA attr-dict `:` type($inA) `to` type($output)";
 }
 
-class MIGraphX_ElementwiseUnaryOp<string name, list<Trait> traits=[]> :
-    MIGraphX_Op<name, traits>,
-    Arguments<(ins AnyMIXRShaped:$inA)>,
-    Results<(outs AnyMIXRShaped:$output)> {
+class MIGraphX_ElementwiseUnaryOp<string name, list<Trait> traits = []>
+    : MIGraphX_Op<name,
+                  !listconcat(traits, [Elementwise,
+                                       AllShapesMatch<["inA", "output"]>])>,
+      Arguments<(ins AnyMIXRShaped:$inA)>,
+      Results<(outs AnyMIXRShaped:$output)> {
   let summary = "Elementwise " # name;
   let assemblyFormat = [{
     $inA attr-dict `:` type($inA) `->` type($output)

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

@@ -334,25 +334,24 @@ LogicalResult ReshapeOp::verify() {
            << 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; });
+  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; });
+  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 ) {
+    // We cannot handle negative dims values that aren't -1
+    if (dimValue < -1) {
       return emitOpError("Non -1 negative values are not supported");
     }
 

mlir/test/Dialect/MIGraphX/invalid.mlir

@@ -1,5 +1,23 @@
 // RUN: rocmlir-opt %s -split-input-file -verify-diagnostics
 
+// ---- migraphx.shaped type ----
+
+// expected-error @+1 {{migraphx.shaped type has 1 elements in its shape but 2 strides defined}}
+func.func @invalid_more_strides_than_shapes(%arg: !migraphx.shaped<1xf32, 1x1>)  {
+  func.return
+}
+
+// -----
+
+// expected-error @+1 {{migraphx.shaped type has 2 elements in its shape but 1 strides defined}}
+func.func @invalid_more_shapes_than_strides(%arg: !migraphx.shaped<1x1xf32, 1>)  {
+  func.return
+}
+
+// -----
+
+// ---- migraphx.reshape ----
+
 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>
@@ -42,6 +60,10 @@ func.func @mlir_neg_one_with_zero(%arg0: !migraphx.shaped<2x4xf16, 0x1>) {
   return
 }
 
+// -----
+
+// ---- migraphx.equal ----
+
 func.func @func_equal(%arg0: !migraphx.shaped<1x36x384x64xi32, 884736x24576x64x1>) -> !migraphx.shaped<1x36x384x64xi32, 884736x24576x64x1> attributes{rock.kernel, rock.arch = ""} {
   %cst = migraphx.literal (dense<1> : tensor<1x36x384x64xi32>) : <1x36x384x64xi32, 884736x24576x64x1>
   %0 = migraphx.add %arg0, %cst : <1x36x384x64xi32, 884736x24576x64x1>, <1x36x384x64xi32, 884736x24576x64x1> -> <1x36x384x64xi32, 884736x24576x64x1>
@@ -50,6 +72,104 @@ func.func @func_equal(%arg0: !migraphx.shaped<1x36x384x64xi32, 884736x24576x64x1
   return %1 : !migraphx.shaped<1x36x384x64xi16, 884736x24576x64x1>
 }
 
+// -----
+
+// ---- migraphx.clip ----
+
+func.func @clip_mismatched_element_types(%arg0: !migraphx.shaped<4x8xf32, 8x1>, %arg1: !migraphx.shaped<4x8xf16, 8x1>, %arg2: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<4x8xf32, 8x1> {
+  // expected-error @+1 {{op failed to verify that all of {x, minVals, maxVals, output} have same element type}}
+  %0 = migraphx.clip %arg0, %arg1, %arg2 : <4x8xf32, 8x1>, <4x8xf16, 8x1>, <4x8xf32, 8x1> -> <4x8xf32, 8x1>
+  return %0 : !migraphx.shaped<4x8xf32, 8x1>
+}
+
+// -----
+
+func.func @clip_mismatched_shapes(%arg0: !migraphx.shaped<4x8xf32, 8x1>, %arg1: !migraphx.shaped<4x4xf32, 4x1>, %arg2: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<4x8xf32, 8x1> {
+  // expected-error @+1 {{op failed to verify that all of {x, minVals, maxVals, output} have same shape}}
+  %0 = migraphx.clip %arg0, %arg1, %arg2 : <4x8xf32, 8x1>, <4x4xf32, 4x1>, <4x8xf32, 8x1> -> <4x8xf32, 8x1>
+  return %0 : !migraphx.shaped<4x8xf32, 8x1>
+}
+
+// -----
+
+// ---- migraphx.where ----
+
+func.func @where_cond_not_bool(%arg0: !migraphx.shaped<4x4xf32, 4x1>, %arg1: !migraphx.shaped<4x4xf32, 4x1>, %arg2: !migraphx.shaped<4x4xf32, 4x1>) -> !migraphx.shaped<4x4xf32, 4x1> {
+  // expected-error @+1 {{'migraphx.where' op operand #0 must be !migraphx.shaped of 8-bit signless integer or 8-bit signed integer or 8-bit unsigned integer values, but got '!migraphx.shaped<4x4xf32, 4x1>'}}
+  %0 = migraphx.where %arg0, %arg1, %arg2 : <4x4xf32, 4x1>, <4x4xf32, 4x1>, <4x4xf32, 4x1> -> <4x4xf32, 4x1>
+  return %0 : !migraphx.shaped<4x4xf32, 4x1>
+}
+
+// -----
+
+func.func @where_mismatched_types(%arg0: !migraphx.shaped<4x4xi8, 4x1>, %arg1: !migraphx.shaped<4x4xf32, 4x1>, %arg2: !migraphx.shaped<4x4xf16, 4x1>) -> !migraphx.shaped<4x4xf32, 4x1> {
+  // expected-error @+1 {{op failed to verify that all of {inA, inB, output} have same element type}}
+  %0 = migraphx.where %arg0, %arg1, %arg2 : <4x4xi8, 4x1>, <4x4xf32, 4x1>, <4x4xf16, 4x1> -> <4x4xf32, 4x1>
+  return %0 : !migraphx.shaped<4x4xf32, 4x1>
+}
+
+// -----
+
+func.func @where_mismatched_shapes(%arg0: !migraphx.shaped<4x4xi8, 4x1>, %arg1: !migraphx.shaped<4x8xf32, 8x1>, %arg2: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<4x8xf32, 8x1> {
+  // expected-error @+1 {{op failed to verify that all of {inA, inB, output, cond} have same shape}}
+  %0 = migraphx.where %arg0, %arg1, %arg2 : <4x4xi8, 4x1>, <4x8xf32, 8x1>, <4x8xf32, 8x1> -> <4x8xf32, 8x1>
+  return %0 : !migraphx.shaped<4x8xf32, 8x1>
+}
+
+// -----
+
+// ---- migraphx.convert ----
+
+func.func @convert_shape_mismatch(%arg0: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<4x4xf16, 4x1> {
+  // expected-error @+1 {{op failed to verify that all of {inA, output} have same shape}}
+  %0 = migraphx.convert %arg0 : <4x8xf32, 8x1> to <4x4xf16, 4x1>
+  return %0 : !migraphx.shaped<4x4xf16, 4x1>
+}
+
+// -----
+
+// ---- migraphx.abs ----
+
+func.func @abs_shape_mismatch(%arg0: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<4x4xf32, 4x1> {
+  // expected-error @+1 {{op failed to verify that all of {inA, output} have same shape}}
+  %0 = migraphx.abs %arg0 : <4x8xf32, 8x1> -> <4x4xf32, 4x1>
+  return %0 : !migraphx.shaped<4x4xf32, 4x1>
+}
+
+// -----
+
+// ---- migraphx.exp ----
+
+func.func @exp_rank_mismatch(%arg0: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<32xf32, 1> {
+  // expected-error @+1 {{op failed to verify that all of {inA, output} have same shape}}
+  %0 = migraphx.exp %arg0 : <4x8xf32, 8x1> -> <32xf32, 1>
+  return %0 : !migraphx.shaped<32xf32, 1>
+}
+
+// -----
+
+// ---- migraphx.relu ----
+
+func.func @relu_shape_mismatch(%arg0: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<2x8xf32, 8x1> {
+  // expected-error @+1 {{op failed to verify that all of {inA, output} have same shape}}
+  %0 = migraphx.relu %arg0 : <4x8xf32, 8x1> -> <2x8xf32, 8x1>
+  return %0 : !migraphx.shaped<2x8xf32, 8x1>
+}
+
+// -----
+
+// ---- migraphx.sigmoid ----
+
+func.func @func_sigmoid_2d_i32(%arg0: !migraphx.shaped<4x8xi32, 8x1>) -> !migraphx.shaped<4x8xi32, 8x1> {
+  // expected-error @+1 {{only support floating point}}
+  %0 = migraphx.sigmoid %arg0 : <4x8xi32, 8x1> -> <4x8xi32, 8x1>
+  return %0 : !migraphx.shaped<4x8xi32, 8x1>
+}
+
+// -----
+
+// ---- migraphx.quant_dot ----
+
 // Test: Only scaleA provided (should fail - both scales required)
 func.func @quant_dot_only_scale_a(
   %arg0: !migraphx.shaped<1x16x512xf4E2M1FN, 8192x512x1>, 
@@ -228,6 +348,8 @@ func.func @migraphx_quant_dot_f4_n_scales(%arg0: !migraphx.shaped<1x16x512xf4E2M
 
 // -----
 
+// ---- migraphx.dot ----
+
 // CHECK-LABEL: func.func @dot_rank_less_than_2
 func.func @dot_rank_less_than_2(%arg0: !migraphx.shaped<320xf16, 1>, %arg1: !migraphx.shaped<320x64xf16, 64x1>) -> !migraphx.shaped<64xf16, 1> {
   // expected-error @+1 {{expect operand to have rank greater or equal to 2}}
@@ -273,51 +395,7 @@ func.func @dot_result_shape_mismatch(%arg0: !migraphx.shaped<2x3x4xf16, 12x4x1>,
 
 // -----
 
-// expected-error @+1 {{migraphx.shaped type has 1 elements in its shape but 2 strides defined}}
-func.func @invalid_more_strides_than_shapes(%arg: !migraphx.shaped<1xf32, 1x1>)  {
-  func.return
-}
-
-// -----
-
-// expected-error @+1 {{migraphx.shaped type has 2 elements in its shape but 1 strides defined}}
-func.func @invalid_more_shapes_than_strides(%arg: !migraphx.shaped<1x1xf32, 1>)  {
-  func.return
-}
-
-// -----
-
-func.func @func_sigmoid_2d_i32(%arg0: !migraphx.shaped<4x8xi32, 8x1>) -> !migraphx.shaped<4x8xi32, 8x1> {
-  // expected-error @+1 {{only support floating point}}
-  %0 = migraphx.sigmoid %arg0 : <4x8xi32, 8x1> -> <4x8xi32, 8x1>
-  return %0 : !migraphx.shaped<4x8xi32, 8x1>
-}
-
-// -----
-
-func.func @where_cond_not_bool(%arg0: !migraphx.shaped<4x4xf32, 4x1>, %arg1: !migraphx.shaped<4x4xf32, 4x1>, %arg2: !migraphx.shaped<4x4xf32, 4x1>) -> !migraphx.shaped<4x4xf32, 4x1> {
-  // expected-error @+1 {{'migraphx.where' op operand #0 must be !migraphx.shaped of 8-bit signless integer or 8-bit signed integer or 8-bit unsigned integer values, but got '!migraphx.shaped<4x4xf32, 4x1>'}}
-  %0 = migraphx.where %arg0, %arg1, %arg2 : <4x4xf32, 4x1>, <4x4xf32, 4x1>, <4x4xf32, 4x1> -> <4x4xf32, 4x1>
-  return %0 : !migraphx.shaped<4x4xf32, 4x1>
-}
-
-// -----
-
-func.func @where_mismatched_types(%arg0: !migraphx.shaped<4x4xi8, 4x1>, %arg1: !migraphx.shaped<4x4xf32, 4x1>, %arg2: !migraphx.shaped<4x4xf16, 4x1>) -> !migraphx.shaped<4x4xf32, 4x1> {
-  // expected-error @+1 {{op failed to verify that all of {inA, inB, output} have same element type}}
-  %0 = migraphx.where %arg0, %arg1, %arg2 : <4x4xi8, 4x1>, <4x4xf32, 4x1>, <4x4xf16, 4x1> -> <4x4xf32, 4x1>
-  return %0 : !migraphx.shaped<4x4xf32, 4x1>
-}
-
-// -----
-
-func.func @where_mismatched_shapes(%arg0: !migraphx.shaped<4x4xi8, 4x1>, %arg1: !migraphx.shaped<4x8xf32, 8x1>, %arg2: !migraphx.shaped<4x8xf32, 8x1>) -> !migraphx.shaped<4x8xf32, 8x1> {
-  // expected-error @+1 {{op failed to verify that all of {inA, inB, output, cond} have same shape}}
-  %0 = migraphx.where %arg0, %arg1, %arg2 : <4x4xi8, 4x1>, <4x8xf32, 8x1>, <4x8xf32, 8x1> -> <4x8xf32, 8x1>
-  return %0 : !migraphx.shaped<4x8xf32, 8x1>
-}
-
-// -----
+// ---- migraphx.slice ----
 
 func.func @invalid_attr_size_mismatch(%input: !migraphx.shaped<10x10xf32, 10x1>) {
   // expected-error @+1 {{op axes, starts, and ends must have the same size}}