`bitcast_convert` of packed constant `U8` (uint8) to "packed" types (`S4`, `S2`, `U4`, `U2`) fails

[janpfeifer]

This following minimal StableHLO code exemplifies it:

• The U8 value 0xE1 if fed as an input parameter and as a constant --> 0x1 nibble is 1 in S4, and 0xE becomes -2 in S4.
• Both values (presumably the same) are %stablehlo.bitcast_convert()'ed to Int4
• And then converted to Int8 for outputting/printing.

But the outputs are different!?

module @TestSubByteDTypes_Int4_AsLiteral {
  func.func @main(%x: tensor<1xui8>) -> (tensor<1x2xi8>, tensor<1x2xi8>, tensor<1xi1>) {
    %0 = "stablehlo.constant"() { value = dense<[225]> : tensor<1xui8> } : () -> tensor<1xui8>
    %1 = "stablehlo.compare"(%0, %x) {
      compare_type = #stablehlo<comparison_type UNSIGNED>,
      comparison_direction = #stablehlo<comparison_direction EQ>
    } : (tensor<1xui8>, tensor<1xui8>) -> tensor<1xi1>
    %2 = "stablehlo.bitcast_convert"(%0) : (tensor<1xui8>) -> tensor<1x2xi4>
    %3 = "stablehlo.convert"(%2) : (tensor<1x2xi4>) -> tensor<1x2xi8>
    %4 = "stablehlo.bitcast_convert"(%x) : (tensor<1xui8>) -> tensor<1x2xi4>
    %5 = "stablehlo.convert"(%4) : (tensor<1x2xi4>) -> tensor<1x2xi8>
    "stablehlo.return"(%3, %5, %1) : (tensor<1x2xi8>, tensor<1x2xi8>, tensor<1xi1>) -> ()
  }
}

If one feeds this function the same value as the constant 0xE1 (225 in decimal) one would expect them to be bitcast to Int4 values [1, -2]. But the outputs of the StableHLO program above are:

1. [[1, 0]] --> Fed as a constant, it ignores the upper nibble (upper 4 bits).
2. [[1, -2]] --> Same value fed as an input parameter, makes the correct stablehlo.bitcast_convert().
3. [true] --> A (in-)sanity check: the input parameter and the constant, both set to [0xE1], compare to true, even though they bitcast_convert to different values.

ps.: Notice PJRT PJRT_Client_BufferFromHostBuffer() method doesn't seem to accept buffers with packed S4 as inputs, so it requires feeding them as U8 and then %stablehlo.bitcast_convert them to S4.

[akuegel]

I suspect this is because Literal does not support packed subbyte elements (https://github.com/openxla/xla/blob/main/xla/literal.h#L1610). When a constant is passed, eventually we will end up constant-folding it which means we will call LiteralBase::BitcastConvert:

xla/xla/literal.cc

Line 1849 in efb7610

absl::StatusOr<Literal> LiteralBase::BitcastConvert(

This does just a memcpy.
I guess for this to work properly, we would need support for packed subbyte elements in Literal.

For your immediate needs, it sounds as if you want to use inputs, anyway, and that works correctly? As a workaround for the broken handling of bitcast_convert from a constant, you could pass the result of bitcast-convert(constant) to XLA.

[janpfeifer]

Not having much understanding of what goes behind XLA, I was expecting something like that, that the literal had their own op (as opposed to becoming values in the execution, like an input). I wonder when does the literal becomes a "normal" value in XLA, and if there is a way to force it ... meaning it only would say as a U8, and the bitcast would use the normal path (same as an input parameter, or the output of a computation) -- again, I have no idea of what's going on inside XLA.

As a work-around now, I'm just handling bitcast(const) by returning an error to the end user (of GoMLX) pointint to this page and asking them to use input instead.

Now I was also not able to feed directly a buffer of packed S4 or U4 -- for some reason the PJRT_Buffer_FromHost call expects one byte per value, hence the need for a bitcast_convert call. The more recent PJRT API has a PJRT_Buffer_Bitcast that works as well, so I can wrap that. Unfortunately, apparently it does a buffer copy :(, I wish it would do it in-place.

[akuegel]

As far as I know, Literal is the data structure that is passed around. So the results you get will be unpacked, the parameters will be unpacked. Why it is like this, I don't know. Most likely because initially we didn't have any packed types in XLA, and then the packing was only implemented for device logic. I will ask around internally.

[janpfeifer]

I can imagine how this came to be like that ... but it seems an inconsistent state: bitcast_convert works differently if the input comes from a parameter (or other computation in the graph) or if it comes from a literal, which is not nice.