TritonToUnstructured: handle BitcastOp on pointer tensors

[sdalvi-quic]

Summary

Handle tt.bitcast on pointer tensors in the TritonToUnstructured pass's offset propagation worklist.

When a kernel bitcasts a tensor of pointers (e.g., tensor<1024x!tt.ptr<i1>>→ tensor<1024x!tt.ptr<i8>>), the pass previously hit the.Default case and failed with "unexpected op in ptr sequence", blocking gather scatter lowering for the entire module.

Approach

The fix adds a .Case<triton::BitcastOp> handler that:

1. Guards against different-stride bitcasts — uses DataLayout::getTypeSizeto compare pointee byte sizes. If they differ (e.g., ptr<i8> → ptr<i16>), the offset cannot be safely reused and the pass emits an error.

2. Reuses the accumulated offset when strides match — when both pointee types have the same byte size (e.g., i1 and i8 both occupy 1 byte in DataLayout), the numeric offset remains valid across the bitcast.

3. Bitcasts the scalar base pointer — creates a new tt.bitcast on the scalar base pointer (kernel arg) so downstream gather/scatter ops see the correct element type.

This is intentionally conservative: only same-stride ptr→ptr bitcasts are supported. Different-stride cases fail explicitly rather than silently producing incorrect code.

Change-Id: I1a7e0eab7e7c9391c8d3a28cfee3a80f7f9d663e

[meta-cla]

Hi @sdalvi-quic!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[sdalvi-quic]

@mitekoth.

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[mitekoth]

LGTM, with minor comments! Will let @nhat-nguyen chime in if this lowering is valid and/if any other changes are required.

[mitekoth]

Please expand the lit test check to be more specific.

[sdalvi-quic]

Added. Thanks!

[mitekoth]

@nhat-nguyen Would be great if you could provide @sdalvi-quic with access to run CI, thanks!

[sdalvi-quic]

Hi @nhat-nguyen, can you please review the PR? Thanks!

[nhat-nguyen]

In general, reusing the previous offset after bitcast is not correct.

I think bitcast is generally hard and ambiguous to handle in this pass. Let's look at this example:

%base_ptr = single ptr of !tt.ptr<i8>
%const_1 = 1
%const_2 = 2
%ptr_0 = tt.addptr %base_ptr %const_1 !tt.ptr<i8> // offset: +1
%ptr_2 = tt.addptr %ptr_0 %const_2 !tt.ptr<i8> // = offset: +3 (since we + 1 above)

%new_ptr = tt.bitcast from !tt.ptr<i8> to !tt.ptr<i16>

%new_ptr_0 = tt.addptr %new_ptr %const_2 // offset: +5 (wrong!)

Each tt.addptr scales the offsets according to the base pointer type (so address + sizeof(sizeof(pointee_type))). When we bitcast, the base has changed so we cannot reuse the old accumulated offsets anymore.

For a scalar pointer, we can fix this by treating the bitcast as the start of a new offset chain (i.e: offset start from 0). The new base pointer will be at the result of the bitcast. We have a patch internally that will handle this. For the above example, this looks roughly like:

// we have accumulated offset = 3 so far, so materialize it using tt.addptr:
%ptr = tt.addptr %base_ptr %const_3
%new_ptr = tt.bitcast from !tt.ptr<i8> to !tt.ptr<i16> // now this is the new base pointer going forward, offset is now reset to 0 too

For tensor of pointers, which is what you are trying to support here, I think there's no good way to represent this generally. This is because after we "materialize" the bitcast, we end up with a tensor of pointers with no single base.

%out_addptr = tt.addptr %out_splat, %offsets : tensor<1024x!tt.ptr<i8>>, tensor<1024xi32>
%bitcast = tt.bitcast %out_addptr from tensor<1024x!tt.ptr<i8>> to tensor<1024x!tt.ptr<i16>>

^ we do not have a single base pointer to transform anymore

We cannot use the original base pointer because the offsets while in !tt.ptr<i8> are scaled by sizeof(i8), while the new offsets in the new type will follow sizeof(new_type).

[sdalvi-quic]

Thanks for the explanation!

The failure we hit is specifically with boolean kernels where Triton widens i1 loads to i8, inserting tt.bitcast on tensor<1024x!tt.ptr<i1>> → tensor<1024x!tt.ptr<i8>>. For this case, both types have effective stride of 1 byte, so reusing the offset is correct.

Two options I see:

1. Handle only same-stride cases in this pass: Add a BitcastOp handler that propagates offsets only when sizeof(src) == sizeof(dst), and errors on other conversions.
2. Eliminate the bitcast upstream: Fix it so the i1→i8 bitcast never reaches this pass (e.g., canonicalize the pointer type earlier in the pipeline). What would be the best place to address this?

Which approach would you recommend?

[nhat-nguyen]

Handle only same-stride cases in this pass: Add a BitcastOp handler that propagates offsets only when sizeof(src) == sizeof(dst), and errors on other conversions.

I'm not sure I follow sizeof(src) == sizeof(dst). We are casting i1 to i8, so sizeof(i1) != sizeof(i8) right?

Do you only need to handle i1 to other larger-width integer types, and that the bitcast is guaranteed not to be used later in other pointer arithmetics?

I think i1 -> larger-width types is safe (like a void* to other types reinterpret_cast). But since this pass just does the transformation generally, we should probably handle this elsewhere.

Eliminate the bitcast upstream: Fix it so the i1→i8 bitcast never reaches this pass (e.g., canonicalize the pointer type earlier in the pipeline). What would be the best place to address this?

Could you elaborate on what the transformation could look like?

[sdalvi-quic]

Handle only same-stride cases in this pass: Add a BitcastOp handler that propagates offsets only when sizeof(src) == sizeof(dst), and errors on other conversions.

By this I wanted to refer to the behavior where MLIR's DataLayout reports
getTypeSize(i1) == 1 and getTypeSize(i8) == 1 (sub-byte types
occupy 1 byte in memory). This would result in same offset for both. I have updated the PR with this logic.

For ptrs with different strides, I updated the approach to guard against those cases and fail the pass as per the existing behavior.

Eliminate the bitcast upstream

About this, I am not sure if there is any more appropriate place for the fix. I am open to suggestions.

[nhat-nguyen]

By this I wanted to refer to the behavior where MLIR's DataLayout reports
getTypeSize(i1) == 1 and getTypeSize(i8) == 1 (sub-byte types
occupy 1 byte in memory).

I'm not sure I follow this. Is this generally true? a (void*) + 1 is different from a (int8*) + 1 even if getTypeSize(i1) == 1 and getTypeSize(i8) == 1. I think this is only safe to do if the bitcast is the last user of an offset chain. If you use the offset afterwards it's still going to be illegal.

[nhat-nguyen]

@sdalvi-quic can you add more tests cases to demonstrate that we're correctly handling multiple bitcast chains?

[sdalvi-quic]

Hi @nhat-nguyen, I have added python and mlir test for multiple bitcast chains.

The primary motivation for this patch is supporting i1 → i8 pointer bitcast, which Triton emits when a kernel operates on boolean tensors. In MLIR's i1 has a store size of 1 byte (no sub-byte packing) similar to i8, so i1 and i8 share the same byte stride. Reusing the accumulated element offset across the bitcast should be fine. The pass rejects cases where the pointee byte size differs (e.g., i8 → i16).

I have also added a Python E2E test (test_bitcast_ptr.py) that verifies correct memory access at runtime. Can you please take a look and let me know if the PR looks good?

[sdalvi-quic]

Hi @nhat-nguyen, I have made the recommended changes. Can you please review the PR and also help with the CI? Thanks!

[nhat-nguyen]

I think this crashes on a loop-carried pointer, though I haven't run it — could you check? The guard only looks at the result type, but then we cast the source type. When the bitcast's source is an scf.for iter-arg, the for-loop handler retypes that iter-arg to the integer offset type in place before this use gets processed, so src.getType() ends up tensor<...xi32> and cast<triton::PointerType> on it asserts. Rough repro:

tt.func public @loop_carried_bitcast(%arg0: !tt.ptr<i1>, %arg1: !tt.ptr<i8>) {
  %c0 = arith.constant 0 : i32
  %c1 = arith.constant 1 : i32
  %c4 = arith.constant 4 : i32
  %step = arith.constant dense<8> : tensor<128xi32>
  %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>
  %1 = tt.splat %arg0 : !tt.ptr<i1> -> tensor<128x!tt.ptr<i1>>
  %2 = tt.addptr %1, %0 : tensor<128x!tt.ptr<i1>>, tensor<128xi32>
  %res = scf.for %i = %c0 to %c4 step %c1 iter_args(%p = %2)
      -> (tensor<128x!tt.ptr<i1>>) : i32 {
    %bc = tt.bitcast %p : tensor<128x!tt.ptr<i1>> -> tensor<128x!tt.ptr<i8>>
    %ld = tt.load %bc : tensor<128x!tt.ptr<i8>>
    %sp = tt.splat %arg1 : !tt.ptr<i8> -> tensor<128x!tt.ptr<i8>>
    %so = tt.addptr %sp, %0 : tensor<128x!tt.ptr<i8>>, tensor<128xi32>
    tt.store %so, %ld : tensor<128x!tt.ptr<i8>>
    %next = tt.addptr %p, %step : tensor<128x!tt.ptr<i1>>, tensor<128xi32>
    scf.yield %next : tensor<128x!tt.ptr<i1>>
  }
  tt.return
}

offsetInfo.ptrType still holds the real pointer type, so reading the pointees off that avoids it (and drops the tensor/scalar dyn_cast branch):

auto offsetInfo = offsetMap.at(src);
Type srcPointee = triton::getPointeeType(offsetInfo.ptrType);
Type dstPointee = triton::getPointeeType(resType);
unsigned srcBytes = dataLayout.getTypeSize(srcPointee);
unsigned dstBytes = dataLayout.getTypeSize(dstPointee);

Would be great to add that loop case as a lit test too.

[sdalvi-quic]

Thanks for the pick. Updated the code.

[nhat-nguyen]

this looks fragile, could we split the i8→i16 rejection into its own // ----- section (or file) with --verify-diagnostics and // expected-error {{bitcast between pointer types with different strides}} right on the op?

[sdalvi-quic]

Yes, added a failing test. Thanks!

[nhat-nguyen]

@sdalvi-quic thanks for adding more tests, generally lgtm. just 2 minor comments and we're good to go :)

[sdalvi-quic]

@nhat-nguyen I have addressed the comments. Can you please review?

[sdalvi-quic]

Hi @nhat-nguyen, can you please review the PR?

[sdalvi-quic]

Ping for review @nhat-nguyen. Thanks!

[nhat-nguyen]

Thanks for the great work!

[nhat-nguyen]

@sdalvi-quic could you fix the failing lit test?

[sdalvi-quic]

@nhat-nguyen, can you please trigger the CI? I fixed the test. Thanks!

[sdalvi-quic]

@nhat-nguyen the build has completed. Can you please merge the PR? Thanks!