Verification of ReduceOp and ReduceWindowOp is not strict enough.

[sdasgup3]

Request description

func.func @variadic_reduce_add(%arg0: tensor<?x?xf32>, %arg1: tensor<?x4xi32>, 
    %arg2: tensor<?x5xi32>, %arg3: tensor<f32>,
    %arg4: tensor<i32>, %arg5: tensor<i32>
    ) -> (tensor<f32>, tensor<i32>, tensor<i32>) {
  %0:3 = mhlo.reduce(%arg0 init: %arg3),
                    (%arg1 init: %arg4),
                    (%arg2 init: %arg5) across dimensions = [0, 1] {someattr}
   : (tensor<?x?xf32>, tensor<?x4xi32>, tensor<?x5xi32>,
      tensor<f32>, tensor<i32>, tensor<i32>
   ) -> (tensor<f32>, tensor<i32>, tensor<i32>)
     reducer(%arg6: tensor<f32>, %arg9: tensor<f32>) 
            (%arg7: tensor<i32>, %arg10: tensor<i32>)
            (%arg8: tensor<i32>, %arg11: tensor<i32>)  {
      %1 = mhlo.add %arg6, %arg9 : tensor<f32>
      %2 = mhlo.add %arg7, %arg10 : tensor<i32>
      %3 = mhlo.add %arg8, %arg11 : tensor<i32>
      mhlo.return %1, %2, %3 : tensor<f32>, tensor<i32>, tensor<i32>
    }
    return %0#0, %0#1, %0#2 : tensor<f32>, tensor<i32>, tensor<i32>
  }

https://source.corp.google.com/piper///depot/google3/third_party/stablehlo/stablehlo/dialect/TypeInference.cpp;rcl=482658831;l=914

input shapes tensor<?x4xi32>, tensor<?x5xi32> are incompatible, but verifier allows it, because it checks the compatibility with the first argument that has type tensor<?x?xf32>.

I think it should infer the "most static" shape first and then try to check whether the arguments are compatible with it. It would be nice to enforce arguments to have the same shape by inserting casts to this inferred type.

[sdasgup3]

It would be nice to enforce arguments to have the same shape by inserting casts to this inferred type.

While re-thinking on that:
StableHLO framework, by itself, will not be adding any extra casts to enforce those constraints. However, "inserting casts" seems like a nice alternative to adding verification checks.

[ghpvnist]

There is a similar issue with ReduceWindowOp

stablehlo/stablehlo/dialect/TypeInference.cpp

Lines 1044 to 1055 in 5e41ad2

	// P2.
	if (!allInputsUnranked) {
	for (uint64_t inputIdx = 0; inputIdx < numInputs; ++inputIdx) {
	if (failed(mlir::verifyCompatibleShape(inputArgTypes[rankedInputIdx],
	inputArgTypes[inputIdx]))) {
	return emitOptionalError(
	location, "expects all inputs to have compatible shapes. Shape at",
	" input-index ", inputIdx,
	" is not compatible with shape at input-index ", rankedInputIdx);
	}
	}
	}

[zhouxin913]

This problem become a little more complicate after we add bounds support in #737: we also need verify bounds are compatible with each other: a.bounds >= b.dim, final bound = min(every bound), etc.f

[sdasgup3]

Adding more ops which involves matching shapes and need to deal with the same issue. As previously discussed it would make sense to collect the size and bound for each dimension individually and use that for verification/type inference.

• ConcatenateOp: shape of operands must match except for the concatenate dimension.
• IfOp/CaseOp: The output shape of the regions should match.
• MapOp: All inputs to have the same shape.
• ScatterOp: All inputs have the same shape.