openxla / stablehlo

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

Improve Tensor and Element prettyprinting #227

Open burmako opened 1 year ago

burmako commented 1 year ago

As discussed, it would be good to align the constant format defined by the spec, and prettyprinting of Tensor and Element. While we're there, let's decide on print and dump (do we need both? do we need to append \n to dumps?).

sdasgup3 commented 1 year ago

Adding one more expectation from pretty-printing: At the moment, tensor are printed following the canonical ordering of elements where each of the lements are printed in its own line. That makes the print of a tensor object pretty verbose.

We might have the following prototype using DenseElementsAttr::print. Below is how a 2x3 shaped tensor with different types would look like:

dense<[[-8, 0, 4], [-6, 7, 5]]> : tensor<2x3xi4>

dense<[[-8.000000e+00, 0.000000e+00, 4.000000e+00], [-6.000000e+00, 7.000000e+00, 5.000000e+00]]> : tensor<2x3xf64>

dense<[[(3.000000e+00,5.000000e+00), (1.500000e+01,1.100000e+01)], [(3.000000e+00,5.000000e+00), (1.500000e+01,1.100000e+01)]]> : tensor<2x2xcomplex<f64>>

With that, the print of tensor, with integer, floating-point, or complex element type, is not only aligned with the spec but also with how input tensors are specified in the test-cases . The later is also important because if the input tensor is following a particular structure, which in this case means explicitly specifying the shape using parenthesis, then it is expected that the output tensor to follow the same.

Also, as suggested by @burmako, we do not have to print above constants entirely in a single line as that wont look elegant in a test-case file with line overflows. Instead we can use CHECK-SAME lit directives to break them.

The downside of the above prototype is that it is not modular: Tensor::print does not depend on Element::print. The Tensor::print collects its individual Element objects are create a DenseElementsAttr out of it to be printed.

print and dump (do we need both? do we need to append \n to dumps?).

In our current implementation, the distinction between dump and print are kind of blurred because we are using Tensor::print in exactly the same way as Tensor::dump uses it.

But having one seems inlined with how MLIR objects handles printing. Maybe in future print can deliver more functionalities (than the current choice of stream handle) parameterized based on arguments and dump can implement a specific flavor of it.

IMO, dump should have the "\n" following what other MLIR object are doing. (This is an after though after I removed it)