[mlir][vector] Relax restriction on reduction distribution

[ThomasRaoux]

Summary

Relax unnecessary restriction when distribution a vector.reduce op. All the float and integer types can be supported by user's lambda.

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
     // Only rank 1 vectors supported.
     if (vectorType.getRank() != 1)
       return rewriter.notifyMatchFailure(
           warpOp, "Only rank 1 reductions can be distributed.");
     // Only warp_size-sized vectors supported.
     if (vectorType.getShape()[0] % warpOp.getWarpSize() != 0)
       return rewriter.notifyMatchFailure(
           warpOp, "Reduction vector dimension must match was size.");
-    // Only f32, i32, f16, i8 element types are supported.
-    if (!reductionOp.getType().isF32() &&
-        !reductionOp.getType().isSignlessInteger(32) &&
-        !reductionOp.getType().isF16() && !reductionOp.getType().isInteger(8))
+    if (!reductionOp.getType().isIntOrFloat())
       return rewriter.notifyMatchFailure(
-          warpOp, "Reduction distribution currently only supports 32bits, f16, "
-                  "and i8 types.");
+          warpOp, "Reduction distribution currently only supports floats and "
+                  "integer types.");
 
     int64_t numElements = vectorType.getShape()[0] / warpOp.getWarpSize();
     // Return vector that will be reduced from the WarpExecuteOnLane0Op.
     unsigned operandIndex = yieldOperand->getOperandNumber();
     SmallVector<Value> yieldValues = {reductionOp.getVector()};
     SmallVector<Type> retTypes = {
         VectorType::get({numElements}, reductionOp.getType())};
     if (reductionOp.getAcc()) {