[mlir][Linalg] Use ConfinedAttr for dimensions of ReduceOp.

[akuegel]

Summary

We can use the new DenseArrayStrictlySorted constraint.

[akuegel]

Rebase.

[akuegel]

Rebase.

[nicolasvasilache]

Nice!

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

def ReduceOp : LinalgStructuredBase_Op<"reduce", [
   }];
 
   let arguments = (ins
     // Input arg
     Variadic<TensorOrMemref>:$inputs,
     // Output arg
     Variadic<TensorOrMemref>:$inits,
 
-    DenseI64ArrayAttr:$dimensions
+    ConfinedAttr<DenseI64ArrayAttr,
+                 [DenseArrayStrictlySorted<DenseI64ArrayAttr>]>:$dimensions
   );
   let results = (outs Variadic<TensorOrMemref>);
   let regions = (region SizedRegion<1>:$combiner);
 
   let extraClassDeclaration = structuredOpsBaseDecls # [{
     // Declare functions necessary for LinalgStructuredInterface.
     ArrayAttr getIteratorTypes();
     ArrayAttr getIndexingMaps();

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (getInits()[i].getType().cast<ShapedType>().getShape() !=
                            << i
                            << " is not equal to the shape at output-index 0.";
     }
   }
   auto inputType = getInputs()[0].getType().cast<ShapedType>();
   auto initType = getInits()[0].getType().cast<ShapedType>();
 
   DenseSet<int64_t> dimensionsToReduce;
-  int64_t lastDimension = -1;
   for (int64_t dimension : dimensionsRef) {
     if (dimension < 0 || dimension >= inputType.getRank()) {
       return emitOpError()
              << "dimensions for reduction should be in the range [0, "
              << inputType.getRank() - 1 << "].";
     }
-    if (dimension <= lastDimension) {
-      return emitOpError()
-             << "reduction dimensions are not in increasing order: "
-             << dimensionsRef;
-    }
-
-    lastDimension = dimension;
     dimensionsToReduce.insert(dimension);
   }
 
   auto inputDims = inputType.getShape();
   auto initDims = initType.getShape();
 
   // Input dimensions that will be left after the reduction.
   SmallVector<int64_t> reducedInputDims;

mlir/test/Dialect/Linalg/invalid.mlir

%reduce = linalg.reduce
       }
   func.return %reduce : tensor<16x64xf32>
 }
 
 // -----
 
 func.func @reduce_duplicate_dimensions(%input: tensor<16x32x64xf32>,
     %init: tensor<16xf32>)  -> tensor<16xf32> {
-  // expected-error @+1 {{'linalg.reduce' op reduction dimensions are not in increasing order: 1, 1}}
+  // expected-error @+1 {{'linalg.reduce' op attribute 'dimensions' failed to satisfy constraint: i64 dense array attribute should be in increasing order}}
   %reduce = linalg.reduce
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<16xf32>)
       dimensions = [1, 1]
       (%in: f32, %out: f32) {
         %0 = arith.addf %in, %out: f32
         linalg.yield %0: f32
       }
   func.return %reduce : tensor<16xf32>
 }
 
 // -----
 
 func.func @reduce_non_increasing_dimensions(%input: tensor<16x32x64xf32>,
     %init: tensor<16xf32>)  -> tensor<16xf32> {
-  // expected-error @+1 {{'linalg.reduce' op reduction dimensions are not in increasing order: 2, 1}}
+  // expected-error @+1 {{'linalg.reduce' op attribute 'dimensions' failed to satisfy constraint: i64 dense array attribute should be in increasing order}}
   %reduce = linalg.reduce
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<16xf32>)
       dimensions = [2, 1]
       (%in: f32, %out: f32) {
         %0 = arith.addf %in, %out: f32
         linalg.yield %0: f32
       }