TritonToUnstructured: handle BitcastOp on pointer tensors

lib/Conversion/TritonToUnstructured/TritonToUnstructuredPass.cpp

@@ -137,10 +137,12 @@
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeRange.h"
@@ -458,6 +460,79 @@ class TritonToUnstructuredPass
                   return success();
                 })
                 .Case<scf::YieldOp>([](auto) { return success(); })
+                .Case<triton::BitcastOp>([&](triton::BitcastOp op) {
+                  auto res = op.getResult();
+                  auto resType = res.getType();
+
+                  if (!triton::isPtrTypeLike(resType)) {
+                    return success();
+                  }
+
+                  // Use the saved ptrType from offsetMap rather than
+                  // src.getType(). When src is an scf.for iter-arg, the
+                  // for-loop handler retypes it to an integer offset type
+                  // in-place before this use is processed, so the live type
+                  // may no longer be a pointer type.
+                  auto src = op.getSrc();
+                  auto offsetInfo = offsetMap.at(src);
+
+                  // Extract pointee types. getPointeeType handles both
+                  // tensor-of-pointers and scalar pointers, returning
+                  // tensor<Nxelem> or elem respectively. We unwrap the
+                  // tensor to get the scalar type for DataLayout.
+                  Type srcPointee = triton::getPointeeType(offsetInfo.ptrType);
+                  Type dstPointee = triton::getPointeeType(resType);
+                  if (auto t = dyn_cast<RankedTensorType>(srcPointee))
+                    srcPointee = t.getElementType();
+                  if (auto t = dyn_cast<RankedTensorType>(dstPointee))
+                    dstPointee = t.getElementType();
+
+                  // Use DataLayout to get the store size in bytes for each
+                  // pointee type. This correctly handles sub-byte types
+                  // (e.g., i1 occupies 1 byte in memory).
+                  auto mod = op->getParentOfType<ModuleOp>();
+                  mlir::DataLayout dataLayout(mod);
+                  unsigned srcBytes = dataLayout.getTypeSize(srcPointee);
+                  unsigned dstBytes = dataLayout.getTypeSize(dstPointee);
+
+                  if (srcBytes != dstBytes) {
+                    op->emitError(
+                        "bitcast between pointer types with different strides "
+                        "is not supported in offset propagation (src size: ")
+                        << srcBytes << " bytes, dst size: " << dstBytes
+                        << " bytes)";
+                    return failure();
+                  }
+
+                  // Safe to reuse offset info — both pointer types have the
+                  // same effective byte stride, so accumulated offsets remain
+                  // valid after the bitcast.
+
+                  // Get the destination pointer type for the base pointer
+                  // bitcast. For tensors, extract the element type (e.g.,
+                  // tensor<128x!tt.ptr<i8>> → !tt.ptr<i8>).
+                  Type dstPtrType;
+                  if (auto resTensorTy = dyn_cast<RankedTensorType>(resType)) {
+                    dstPtrType = resTensorTy.getElementType();
+                  } else {
+                    dstPtrType = resType;
+                  }
+
+                  // Bitcast the base pointer to match the new pointee type.
+                  OpBuilder b{op};
+                  Value newBasePtr = triton::BitcastOp::create(
+                      b, op->getLoc(), dstPtrType, offsetInfo.ptr);
+
+                  PtrOffset newOffsetInfo{newBasePtr, resType,
+                                          offsetInfo.bitWidth,
+                                          offsetInfo.offset};
+
+                  offsetMap.insert({res, newOffsetInfo});
+                  workList.push(res);
+                  toDelete.push_back(op);
+
+                  return success();
+                })
                 .Case<triton::CatOp>([](triton::CatOp op) {
                   op->emitError("Do not support gather / scatter with multiple "
                                 "bases yet");
@@ -584,11 +659,23 @@ class TritonToUnstructuredPass
   }
 
   void runOnOperation() override {
-    if (failed(processUnstructuredPtrs(offsetBitWidth))) {
-      getOperation()->emitWarning(
-          "Cannot transform tensor of pointers into a single base pointer "
-          "with tensor of offsets");
-      return;
+    bool emittedError = false;
+    {
+      mlir::ScopedDiagnosticHandler diagHandler(
+          &getContext(), [&](mlir::Diagnostic &diag) {
+            if (diag.getSeverity() == mlir::DiagnosticSeverity::Error)
+              emittedError = true;
+            return failure();
+          });
+
+      if (failed(processUnstructuredPtrs(offsetBitWidth))) {
+        if (!emittedError) {
+          getOperation()->emitWarning(
+              "Cannot transform tensor of pointers into a single base "
+              "pointer with tensor of offsets");
+        }
+        return;
+      }
     }
 
     PassManager pm(&getContext(), getOperation().getOperationName());

python/examples/test_bitcast_ptr.py

@@ -0,0 +1,71 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates, Microsoft Corporation.
+# Licensed under the MIT license.
+
+# Test for bitcast pointer handling in TritonToUnstructuredPass.
+# Verifies that pointer arithmetic and bitcasts between same-size types
+# (i1/i8 — both 1 byte in MLIR DataLayout) access the correct memory
+# addresses via the unstructured (gather/scatter) lowering path.
+
+import pytest
+import torch
+
+import triton
+import triton.language as tl
+
+from triton.backends.triton_shared.driver import CPUDriver
+
+
+@triton.jit
+def bitcast_ptr_kernel(input_ptr, output_ptr, BLOCK: tl.constexpr):
+    idx = tl.arange(0, BLOCK)
+    ptr_i1 = input_ptr + 64 + idx               # addptr before bitcast
+    ptr_i8 = ptr_i1.to(tl.pointer_type(tl.int8), bitcast=True)  # bitcast i1->i8
+    ptr_final = ptr_i8 + 32                      # addptr after bitcast
+    tl.store(output_ptr + idx, tl.load(ptr_final))
+
+
+def test_bitcast_ptr(device):
+    """addptr -> bitcast(i1->i8) -> addptr -> load. Reads at base+96+idx."""
+    if device == "cpu":
+        triton.runtime.driver.set_active(CPUDriver())
+
+    input_buf = torch.arange(0, 256, device=device, dtype=torch.uint8)
+    output_buf = torch.zeros(128, device=device, dtype=torch.uint8)
+
+    bitcast_ptr_kernel[(1,)](
+        input_buf.view(torch.bool), output_buf, BLOCK=128
+    )
+
+    # Kernel reads at base + 64 + idx + 32 = base + 96 + idx
+    expected = input_buf[96:224]
+    assert torch.equal(output_buf, expected)
+
+
+@triton.jit
+def bitcast_chain_kernel(input_ptr, output_ptr, BLOCK: tl.constexpr):
+    idx = tl.arange(0, BLOCK)
+    ptr_i1 = input_ptr + 16 + idx                # addptr on ptr<i1>
+    ptr_i8 = ptr_i1.to(tl.pointer_type(tl.int8), bitcast=True)  # bitcast i1->i8
+    ptr_i8_2 = ptr_i8 + 8                        # addptr on ptr<i8>
+    ptr_i1_2 = ptr_i8_2.to(tl.pointer_type(tl.int1), bitcast=True)  # bitcast i8->i1
+    ptr_i1_3 = ptr_i1_2 + 4                      # addptr on ptr<i1>
+    ptr_i8_3 = ptr_i1_3.to(tl.pointer_type(tl.int8), bitcast=True)  # bitcast i1->i8
+    tl.store(output_ptr + idx, tl.load(ptr_i8_3))
+
+
+def test_bitcast_chain(device):
+    """addptr -> bitcast(i1->i8) -> addptr -> bitcast(i8->i1) -> addptr -> bitcast(i1->i8) -> load.
+    Multiple bitcast chain. Reads at base+16+8+4 = base+28+idx."""
+    if device == "cpu":
+        triton.runtime.driver.set_active(CPUDriver())
+
+    input_buf = torch.arange(0, 256, device=device, dtype=torch.uint8)
+    output_buf = torch.zeros(64, device=device, dtype=torch.uint8)
+
+    bitcast_chain_kernel[(1,)](
+        input_buf.view(torch.bool), output_buf, BLOCK=64
+    )
+
+    # Kernel reads at base + 16 + idx + 8 + 4 = base + 28 + idx
+    expected = input_buf[28:92]
+    assert torch.equal(output_buf, expected)

test/Conversion/TritonToUnstructured/bitcast_ptr_unstructured.mlir

@@ -0,0 +1,124 @@
+// RUN: triton-shared-opt --triton-to-unstructured --split-input-file %s | FileCheck %s
+
+// Test 1: addptr -> bitcast(i1->i8) -> load
+// sizeof(i1)==sizeof(i8)==1 byte, so offset N equals N bytes for both types.
+
+module {
+  tt.func public @bitcast_ptr_to_ptr(%arg0: !tt.ptr<i1>, %arg1: !tt.ptr<i8>) {
+    %0 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32>
+    %1 = tt.splat %arg0 : !tt.ptr<i1> -> tensor<1024x!tt.ptr<i1>>
+    %2 = tt.addptr %1, %0 : tensor<1024x!tt.ptr<i1>>, tensor<1024xi32>
+    %3 = tt.bitcast %2 : tensor<1024x!tt.ptr<i1>> -> tensor<1024x!tt.ptr<i8>>
+    %4 = tt.load %3 : tensor<1024x!tt.ptr<i8>>
+    %5 = tt.splat %arg1 : !tt.ptr<i8> -> tensor<1024x!tt.ptr<i8>>
+    %6 = tt.addptr %5, %0 : tensor<1024x!tt.ptr<i8>>, tensor<1024xi32>
+    tt.store %6, %4 : tensor<1024x!tt.ptr<i8>>
+    tt.return
+  }
+}
+
+// CHECK-LABEL: tt.func public @bitcast_ptr_to_ptr
+// CHECK: [[base:%.+]] = tt.bitcast %arg0 : !tt.ptr<i1> -> !tt.ptr<i8>
+// CHECK: tts.gather [[base]][{{.+}}] : (<i8>, tensor<1024xi32>) -> tensor<1024xi8>
+// CHECK: tts.scatter
+
+// -----
+
+// Test 2: addptr -> bitcast(i1->i8) -> addptr -> load/store
+// Both offsets accumulate via arith.addi since sizeof(i1)==sizeof(i8)==1.
+// Proves addptr after bitcast produces correct combined offset.
+
+module {
+  tt.func public @addptr_bitcast_i1_i8_addptr(%arg0: !tt.ptr<i1>, %arg1: !tt.ptr<i8>) {
+    %cst = arith.constant dense<32> : 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>
+    %3 = tt.bitcast %2 : tensor<128x!tt.ptr<i1>> -> tensor<128x!tt.ptr<i8>>
+    %4 = tt.addptr %3, %cst : tensor<128x!tt.ptr<i8>>, tensor<128xi32>
+    %5 = tt.load %4 : tensor<128x!tt.ptr<i8>>
+    %6 = tt.splat %arg1 : !tt.ptr<i8> -> tensor<128x!tt.ptr<i8>>
+    %7 = tt.addptr %6, %0 : tensor<128x!tt.ptr<i8>>, tensor<128xi32>
+    tt.store %7, %5 : tensor<128x!tt.ptr<i8>>
+    tt.return
+  }
+}
+
+// CHECK-LABEL: tt.func public @addptr_bitcast_i1_i8_addptr
+// CHECK: [[BASE2:%.+]] = tt.bitcast %arg0 : !tt.ptr<i1> -> !tt.ptr<i8>
+// CHECK: [[OFF:%.+]] = arith.addi {{.+}}, {{.+}} : tensor<128xi32>
+// CHECK: tts.gather [[BASE2]]{{\[}}[[OFF]]{{\]}} : (<i8>, tensor<128xi32>) -> tensor<128xi8>
+// CHECK: tts.scatter
+
+// -----
+
+// Test 3: addptr -> bitcast(i1->i8) -> addptr -> bitcast(i8->i1) -> addptr -> bitcast(i1->i8) -> load
+// Multiple bitcast chain (i1->i8->i1->i8). All types are 1-byte stride,
+// so offsets accumulate correctly through all bitcasts.
+// Final offset = range + 32 + 16, base ends as ptr<i8>.
+
+module {
+  tt.func public @bitcast_i1_i8_i1_chain(%arg0: !tt.ptr<i1>, %arg1: !tt.ptr<i8>) {
+    %cst = arith.constant dense<32> : tensor<128xi32>
+    %cst_0 = arith.constant dense<16> : 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>
+    %3 = tt.bitcast %2 : tensor<128x!tt.ptr<i1>> -> tensor<128x!tt.ptr<i8>>
+    %4 = tt.addptr %3, %cst : tensor<128x!tt.ptr<i8>>, tensor<128xi32>
+    %5 = tt.bitcast %4 : tensor<128x!tt.ptr<i8>> -> tensor<128x!tt.ptr<i1>>
+    %6 = tt.addptr %5, %cst_0 : tensor<128x!tt.ptr<i1>>, tensor<128xi32>
+    %7 = tt.bitcast %6 : tensor<128x!tt.ptr<i1>> -> tensor<128x!tt.ptr<i8>>
+    %8 = tt.load %7 : tensor<128x!tt.ptr<i8>>
+    %9 = tt.splat %arg1 : !tt.ptr<i8> -> tensor<128x!tt.ptr<i8>>
+    %10 = tt.addptr %9, %0 : tensor<128x!tt.ptr<i8>>, tensor<128xi32>
+    tt.store %10, %8 : tensor<128x!tt.ptr<i8>>
+    tt.return
+  }
+}
+
+// CHECK-LABEL: tt.func public @bitcast_i1_i8_i1_chain
+// CHECK: tt.bitcast %arg0 : !tt.ptr<i1> -> !tt.ptr<i8>
+// CHECK: [[OFF1:%.+]] = arith.addi {{.+}}, {{.+}} : tensor<128xi32>
+// CHECK: tt.bitcast {{.+}} : !tt.ptr<i8> -> !tt.ptr<i1>
+// CHECK: [[OFF2:%.+]] = arith.addi [[OFF1]], {{.+}} : tensor<128xi32>
+// CHECK: [[BC3:%.+]] = tt.bitcast {{.+}} : !tt.ptr<i1> -> !tt.ptr<i8>
+// CHECK: tts.gather [[BC3]]{{\[}}[[OFF2]]{{\]}} : (<i8>, tensor<128xi32>) -> tensor<128xi8>
+// CHECK: tts.scatter
+
+// -----
+
+// Test 4: bitcast inside scf.for with loop-carried pointer iter-arg.
+// The scf.for handler retypes the iter-arg to an integer offset type before
+// the bitcast handler processes it. The bitcast handler must use the saved
+// ptrType from offsetMap (not src.getType()) to avoid asserting on the
+// already-retyped integer type.
+
+module {
+  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
+  }
+}
+
+// CHECK-LABEL: tt.func public @loop_carried_bitcast
+// CHECK: scf.for
+// CHECK:   [[BC:%.+]] = tt.bitcast %arg0 : !tt.ptr<i1> -> !tt.ptr<i8>
+// CHECK:   tts.gather [[BC]][{{.+}}] : (<i8>, tensor<128xi32>) -> tensor<128xi8>
+// CHECK:   tts.scatter

test/Conversion/TritonToUnstructured/bitcast_ptr_unstructured_invalid.mlir

@@ -0,0 +1,27 @@
+// RUN: triton-shared-opt --triton-to-unstructured --split-input-file --verify-diagnostics %s
+
+// Test: bitcast between pointer types with different pointee byte sizes is rejected.
+// The first bitcast (i1->i8) is valid (both 1 byte), but the second bitcast
+// (i8->i16) has different pointee byte sizes (1 byte vs 2 bytes). The pass
+// rejects this because the accumulated element-offset would be misinterpreted:
+// offset N means N*1 bytes for ptr<i8> but N*2 bytes for ptr<i16>.
+
+module {
+  tt.func public @bitcast_i8_to_i16_rejected(%arg0: !tt.ptr<i1>, %arg1: !tt.ptr<i16>) {
+    %cst = arith.constant dense<32> : tensor<128xi32>
+    %cst_1 = arith.constant dense<4> : 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>
+    %3 = tt.bitcast %2 : tensor<128x!tt.ptr<i1>> -> tensor<128x!tt.ptr<i8>>
+    %4 = tt.addptr %3, %cst : tensor<128x!tt.ptr<i8>>, tensor<128xi32>
+    // expected-error @+1 {{bitcast between pointer types with different strides}}
+    %5 = tt.bitcast %4 : tensor<128x!tt.ptr<i8>> -> tensor<128x!tt.ptr<i16>>
+    %6 = tt.addptr %5, %cst_1 : tensor<128x!tt.ptr<i16>>, tensor<128xi32>
+    %7 = tt.load %6 : tensor<128x!tt.ptr<i16>>
+    %8 = tt.splat %arg1 : !tt.ptr<i16> -> tensor<128x!tt.ptr<i16>>
+    %9 = tt.addptr %8, %0 : tensor<128x!tt.ptr<i16>>, tensor<128xi32>
+    tt.store %9, %7 : tensor<128x!tt.ptr<i16>>
+    tt.return
+  }
+}