openxla / stablehlo

Backward compatible ML compute opset inspired by HLO/MHLO
Apache License 2.0
398 stars 109 forks source link

Create inferFooOp functions for some infeasible ops #867

Open burmako opened 1 year ago

burmako commented 1 year ago

At the moment, status.md has some ops marked as infeasible meaning that we cannot fully infer their return types from their operands and attributes. These ops can be categorized as follows: A) Ops where result's element type is load-bearing: bitcast_convert, convert, convolution, dot_general, uniform_quantize. B) Static ops where result's shape is load-bearing: broadcast_in_dim, iota, infeed, recv, reshape, rng_bit_generator. C) Dynamic ops where result's shape is passed as operands: dynamic_broadcast_in_dim, dynamic_conv, dynamic_gather, dynamic_iota, dynamic_pad, dynamic_reshape, real_dynamic_slice, rng. D) The entire return type is load bearing: custom_call.

To support #622, it will be useful to create inferFooOp functions in TypeInference.h for ops from groups A-C. These functions will return SmallVectorImpl<ShapedTypeComponents>& and only fill in the shape part. Even though these functions won't be usable from the existing type inference framework (missing element types will mean that type inference will fail), they will allow sharing shape inference logic with #622.

zhouxin913 commented 1 year ago

Let me be more specific of what will be done here: This will introduce InferFooOp as a util function in TypeInference.h instead of involving shape inference interface in the op level, which is FooOp::inferReturnTypeComponents(). Thus, we still rely on the current logic in verifiers to check the inferred type VS. real result type.

So a plan will be:

  1. Move most code from existing verifyFooOp to the new inferFooOp (when they are hard to split).
  2. Keep current verifyFooOp interface unchanged, but change the body to: 2.1 Call inferFooOp 2.2 Then check the inferred type VS. real result type
  3. Do we need a special name for these new methods? Like InferFooOpPartially() or a special comment to make it clear that this method is different from other complete infer methods?
burmako commented 1 year ago

Thank you for going into details about the plan! Yes, this is what I was planning to do.

"Do we need a special name for these new methods?". I would propose that we use the same name. After all, the type signature of these functions will be the same as the type signature of the regular inferReturnTypeComponents-style functions. The fact that the current version of the builder generator infrastructure doesn't work well with these methods is an implementation detail.

burmako commented 1 year ago

Category A is handled in #898. Category B is handled is in #891. Category C is still in progress.

burmako commented 1 year ago

Unassigning in favor of higher-priority work. Having a uniform interface to shape refinement will be useful for #1160, but in the meanwhile the existing adhoc implementations of shape refinement patterns are working just fine for #622.