diff --git a/include/flydsl/Dialect/Fly/Transforms/MemrefLowering.td b/include/flydsl/Dialect/Fly/Transforms/MemrefLowering.td index e73e69480b..25f21ca869 100644 --- a/include/flydsl/Dialect/Fly/Transforms/MemrefLowering.td +++ b/include/flydsl/Dialect/Fly/Transforms/MemrefLowering.td @@ -12,8 +12,6 @@ def : Pat<(Fly_PtrLoadOp Fly_IntTuple:$int_tuple), def : Pat<(Fly_AddOffsetOp Fly_IntTuple:$int_tuple, Fly_IntTuple:$offset), (Fly_IntTupleAddOp $int_tuple, $offset)>; -def : Pat<(Fly_AddOffsetOp (Fly_AddOffsetOp Fly_Pointer:$ptr, Fly_IntTuple:$offset1), Fly_IntTuple:$offset2), - (Fly_AddOffsetOp $ptr, (Fly_IntTupleAddOp $offset1, $offset2))>; def : Pat<(Fly_DecompositionOp Fly_MemRef:$memref), (replaceWithValue $memref), diff --git a/lib/Dialect/Fly/Transforms/LayoutLowering.cpp b/lib/Dialect/Fly/Transforms/LayoutLowering.cpp index 1ad4842e5e..44c7b063b2 100644 --- a/lib/Dialect/Fly/Transforms/LayoutLowering.cpp +++ b/lib/Dialect/Fly/Transforms/LayoutLowering.cpp @@ -2893,6 +2893,66 @@ class MemRefStoreVecOpLowering : public OpRewritePattern { } }; +/// Canonicalize an add_offset chain to `add_offset(add_offset(ptr, dyn), static)` +/// -- at most one dynamic layer plus one static tail. +/// +/// Fusing a runtime offset with a compile-time constant makes every access +/// compute its own address off the runtime base, so the backend materializes +/// one base register per access instead of sharing one. Keeping the static +/// part as a separate outer offset preserves the shared base. +/// +/// Convergence: each rewrite either drops a level from the chain (fuse) or +/// moves the static offset strictly outward (swap), and the resulting +/// `dyn -> static` form is rejected by the already-canonical check, so the +/// pattern cannot loop. +class AddOffsetCanonicalization : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(AddOffsetOp op, PatternRewriter &rewriter) const override { + auto inner = op.getPtr().getDefiningOp(); + if (!inner) + return failure(); + // Rewrite pointer based add_offset only. IntTuple bases are handled by the int_tuple_add rule. + if (!isa(inner.getPtr().getType())) + return failure(); + auto innerTy = dyn_cast(inner.getOffset().getType()); + auto outerTy = dyn_cast(op.getOffset().getType()); + if (!innerTy || !outerTy) + return failure(); + if (!innerTy.getAttr().isLeafInt() || !outerTy.getAttr().isLeafInt()) + return failure(); + + bool innerStatic = innerTy.getAttr().getLeafAsInt().isStatic(); + bool outerStatic = outerTy.getAttr().getLeafAsInt().isStatic(); + + // Already canonical. + if (!innerStatic && outerStatic) + return failure(); + + Location loc = op.getLoc(); + + // static -> dynamic: swap so the static offset becomes the outer tail. + if (innerStatic && !outerStatic) { + // The swap rebuilds the inner op, so with other users the old one + // survives and we would add an add_offset instead of replacing one. + // Fusion below has no such problem: it rewrites only the outer op, so + // the inner one stays valid for its other users. + if (!inner->hasOneUse()) + return failure(); + Value dynBase = AddOffsetOp::create(rewriter, loc, inner.getPtr(), op.getOffset()); + rewriter.replaceOpWithNewOp(op, dynBase, inner.getOffset()); + return success(); + } + + // static -> static and dynamic -> dynamic both fuse: no runtime value is + // merged with a constant. + Value sum = IntTupleAddOp::create(rewriter, loc, inner.getOffset(), op.getOffset()); + rewriter.replaceOpWithNewOp(op, inner.getPtr(), sum); + return success(); + } +}; + class MemRefAllocaOpLowering : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; @@ -2997,6 +3057,7 @@ class FlyLayoutLoweringPass // MemRef/Ptr operations patterns.add(context); patterns.add(context); + patterns.add(context); // Utility ops patterns.add(context); diff --git a/tests/mlir/Conversion/nested_add_offset.mlir b/tests/mlir/Conversion/nested_add_offset.mlir new file mode 100644 index 0000000000..7a31d2d9ad --- /dev/null +++ b/tests/mlir/Conversion/nested_add_offset.mlir @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2025 FlyDSL Project Contributors +// RUN: %fly-opt %s --fly-layout-lowering --convert-fly-to-rocdl | FileCheck %s + +// Conversion-level guard for the add_offset canonicalization (issue #898). +// +// A runtime offset fused with a compile-time constant makes every access +// compute its own address, so the backend materializes one base register per +// access. Keeping `add_offset(add_offset(ptr, dyn), static)` split lets a +// single runtime GEP feed many constant GEPs. +// +// Structural assertions only: no ISA register numbers are pinned here. + +// === LDS: one runtime GEP shared by two constant tails === +// +// allocBytes/allocAlign make this lower to a real relocatable @__shared_alloc_* +// LDS symbol, which is the shape the issue reproduces with. +// CHECK-LABEL: gpu.func @lds_shared_runtime_base +gpu.module @m_lds { + gpu.func @lds_shared_runtime_base(%x: i32) kernel { + %smem = fly.make_ptr() {dictAttrs = {allocBytes = 4096 : i64, allocAlign = 128 : i64}} : () -> !fly.ptr + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %c2 = fly.make_int_tuple() : () -> !fly.int_tuple<128> + // The runtime index must reach the GEP unmodified: no `addi %x, const`. + // CHECK: %[[SYM:.*]] = llvm.mlir.addressof @__shared_alloc_{{[0-9]+}} : !llvm.ptr<3> + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = llvm.getelementptr %[[SYM]][%{{.*}}] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, bf16 + // CHECK: %[[C1:.*]] = arith.constant 64 : i32 + // CHECK: %[[PA:.*]] = llvm.getelementptr %[[BASE]][%[[C1]]] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, bf16 + // CHECK: %[[C2:.*]] = arith.constant 128 : i32 + // CHECK: %[[PB:.*]] = llvm.getelementptr %[[BASE]][%[[C2]]] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, bf16 + // CHECK: llvm.load %[[PA]] + // CHECK: llvm.load %[[PB]] + %p0 = fly.add_offset(%smem, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %pa = fly.add_offset(%p0, %c1) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %pb = fly.add_offset(%p0, %c2) : (!fly.ptr, !fly.int_tuple<128>) -> !fly.ptr + %a = fly.ptr.load(%pa) : (!fly.ptr) -> bf16 + %b = fly.ptr.load(%pb) : (!fly.ptr) -> bf16 + gpu.return + } +} + +// ----- + +// === Swizzled LDS: applicability boundary === +// +// A non-trivial swizzle is applied at the load, to the FINAL address, so +// XOR(base + c) != XOR(base) + c and the shared-base shape does not survive. +// The canonicalization is neutral here, not harmful: the runtime base is still +// a single GEP and no runtime+constant fusion happens. This test pins that +// "neutral, still correct" behaviour so the boundary stays visible. +// CHECK-LABEL: gpu.func @lds_swizzled_neutral +gpu.module @m_swz { + gpu.func @lds_swizzled_neutral(%x: i32) kernel { + %smem = fly.make_ptr() {dictAttrs = {allocBytes = 4096 : i64, allocAlign = 128 : i64}} : () -> !fly.ptr> + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + // CHECK: %[[SYM:.*]] = llvm.mlir.addressof @__shared_alloc_{{[0-9]+}} : !llvm.ptr<3> + // CHECK: %[[BASE:.*]] = llvm.getelementptr %[[SYM]][%{{.*}}] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, bf16 + // CHECK: %[[C1:.*]] = arith.constant 64 : i32 + // CHECK: %[[PA:.*]] = llvm.getelementptr %[[BASE]][%[[C1]]] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, bf16 + // The swizzle consumes the whole address, constant tail included. + // CHECK: %[[I:.*]] = llvm.ptrtoint %[[PA]] : !llvm.ptr<3> to i64 + // CHECK: arith.xori %[[I]] + %p0 = fly.add_offset(%smem, %d) : (!fly.ptr>, !fly.int_tuple) -> !fly.ptr> + %pa = fly.add_offset(%p0, %c1) : (!fly.ptr>, !fly.int_tuple<64>) -> !fly.ptr> + %a = fly.ptr.load(%pa) : (!fly.ptr>) -> bf16 + gpu.return + } +} + +// ----- + +// === Buffer descriptor: separate offset adds, no fusion === +// +// Buffer pointers add into the fat-pointer offset field instead of emitting a +// GEP. The runtime and the constant must stay two separate adds, so one +// runtime offset can be shared by several constant tails. +// CHECK-LABEL: @buffer_runtime_then_static +func.func @buffer_runtime_then_static(%ptr: !fly.ptr, %x: i32) -> f32 { + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + // Runtime offset added on its own... + // CHECK: %[[O0:.*]] = llvm.extractvalue %{{.*}}[1] + // CHECK: arith.addi %[[O0]], %{{.*}} : i32 + // ...then the constant added separately, never merged into the runtime add. + // CHECK: %[[C1:.*]] = arith.constant 64 : i32 + // CHECK: %[[O1:.*]] = llvm.extractvalue %{{.*}}[1] + // CHECK: arith.addi %[[O1]], %[[C1]] : i32 + %p0 = fly.add_offset(%ptr, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %pa = fly.add_offset(%p0, %c1) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %a = fly.ptr.load(%pa) : (!fly.ptr) -> f32 + return %a : f32 +} diff --git a/tests/mlir/Transforms/layout_lowering.mlir b/tests/mlir/Transforms/layout_lowering.mlir index 82e8135f17..e9b2045b17 100644 --- a/tests/mlir/Transforms/layout_lowering.mlir +++ b/tests/mlir/Transforms/layout_lowering.mlir @@ -446,3 +446,182 @@ func.func @test_load_vec_coord_tensor() -> vector<6xi32> { %vec = fly.memref.load_vec(%ct) : (!fly.coord_tensor<0, (2, 3) : (0, 1)>) -> vector<6xi32> return %vec : vector<6xi32> } + +// ----- + +// === add_offset canonicalization (issue #898) === +// +// Normal form is `add_offset(add_offset(ptr, dyn), static)`: at most one dynamic +// layer plus one static tail. Fusing a runtime offset with a compile-time +// constant would give every access its own base register instead of sharing one. + +// dynamic -> static is already canonical: the two levels must stay split so +// many accesses share one runtime base instead of each computing its own. +// CHECK-LABEL: @test_add_offset_dyn_static_kept +func.func @test_add_offset_dyn_static_kept(%ptr: !fly.ptr, %x: i32) -> f32 { + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c = fly.make_int_tuple() : () -> !fly.int_tuple<64> + // The runtime offset must never be merged into an arith.addi with the constant. + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: fly.add_offset(%[[BASE]], %{{.*}}) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p1 = fly.add_offset(%ptr, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %p2 = fly.add_offset(%p1, %c) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %v = fly.ptr.load(%p2) : (!fly.ptr) -> f32 + return %v : f32 +} + +// static -> dynamic is reordered into the canonical dynamic -> static form. +// CHECK-LABEL: @test_add_offset_static_dyn_swapped +func.func @test_add_offset_static_dyn_swapped(%ptr: !fly.ptr, %x: i32) -> f32 { + %c = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: fly.add_offset(%[[BASE]], %{{.*}}) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p1 = fly.add_offset(%ptr, %c) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p2 = fly.add_offset(%p1, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %v = fly.ptr.load(%p2) : (!fly.ptr) -> f32 + return %v : f32 +} + +// static -> static still fuses: no runtime value is merged with a constant. +// CHECK-LABEL: @test_add_offset_static_static_fused +func.func @test_add_offset_static_static_fused(%ptr: !fly.ptr) -> f32 { + %a = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %b = fly.make_int_tuple() : () -> !fly.int_tuple<32> + // CHECK: %[[C:.*]] = fly.make_int_tuple() : () -> !fly.int_tuple<96> + // CHECK: fly.add_offset(%{{.*}}, %[[C]]) : (!fly.ptr, !fly.int_tuple<96>) -> !fly.ptr + // CHECK-NOT: fly.add_offset + %p1 = fly.add_offset(%ptr, %a) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p2 = fly.add_offset(%p1, %b) : (!fly.ptr, !fly.int_tuple<32>) -> !fly.ptr + %v = fly.ptr.load(%p2) : (!fly.ptr) -> f32 + return %v : f32 +} + +// dynamic -> dynamic still fuses: no constant is lost into a runtime value. +// CHECK-LABEL: @test_add_offset_dyn_dyn_fused +func.func @test_add_offset_dyn_dyn_fused(%ptr: !fly.ptr, %x: i32, %y: i32) -> f32 { + %a = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %b = fly.make_int_tuple(%y) : (i32) -> !fly.int_tuple + // CHECK: %[[SUM:.*]] = arith.addi + // CHECK: %[[T:.*]] = fly.make_int_tuple(%[[SUM]]) + // CHECK: fly.add_offset(%{{.*}}, %[[T]]) + // CHECK-NOT: fly.add_offset + %p1 = fly.add_offset(%ptr, %a) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %p2 = fly.add_offset(%p1, %b) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %v = fly.ptr.load(%p2) : (!fly.ptr) -> f32 + return %v : f32 +} + +// Two constants behind one runtime offset collapse into a single static tail. +// CHECK-LABEL: @test_add_offset_dyn_c1_c2 +func.func @test_add_offset_dyn_c1_c2(%ptr: !fly.ptr, %x: i32) -> f32 { + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %c2 = fly.make_int_tuple() : () -> !fly.int_tuple<32> + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: %[[C:.*]] = fly.make_int_tuple() : () -> !fly.int_tuple<96> + // CHECK: fly.add_offset(%[[BASE]], %[[C]]) : (!fly.ptr, !fly.int_tuple<96>) -> !fly.ptr + %p1 = fly.add_offset(%ptr, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %p2 = fly.add_offset(%p1, %c1) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p3 = fly.add_offset(%p2, %c2) : (!fly.ptr, !fly.int_tuple<32>) -> !fly.ptr + %v = fly.ptr.load(%p3) : (!fly.ptr) -> f32 + return %v : f32 +} + +// A three-level static -> dynamic -> static chain converges to the normal form +// with both constants merged; the pattern must not ping-pong (pass must not fail). +// CHECK-LABEL: @test_add_offset_static_dyn_static +func.func @test_add_offset_static_dyn_static(%ptr: !fly.ptr, %x: i32) -> f32 { + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c2 = fly.make_int_tuple() : () -> !fly.int_tuple<32> + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: %[[C:.*]] = fly.make_int_tuple() : () -> !fly.int_tuple<96> + // CHECK: fly.add_offset(%[[BASE]], %[[C]]) : (!fly.ptr, !fly.int_tuple<96>) -> !fly.ptr + %p1 = fly.add_offset(%ptr, %c1) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p2 = fly.add_offset(%p1, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %p3 = fly.add_offset(%p2, %c2) : (!fly.ptr, !fly.int_tuple<32>) -> !fly.ptr + %v = fly.ptr.load(%p3) : (!fly.ptr) -> f32 + return %v : f32 +} + +// One runtime base shared by two different static tails: the base must stay a +// single add_offset, which is what collapses the ds_read base registers. +// CHECK-LABEL: @test_add_offset_shared_runtime_base +func.func @test_add_offset_shared_runtime_base(%ptr: !fly.ptr, %x: i32) -> f32 { + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %c2 = fly.make_int_tuple() : () -> !fly.int_tuple<128> + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: fly.add_offset(%[[BASE]], %{{.*}}) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + // CHECK: fly.add_offset(%[[BASE]], %{{.*}}) : (!fly.ptr, !fly.int_tuple<128>) -> !fly.ptr + %p0 = fly.add_offset(%ptr, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %pa = fly.add_offset(%p0, %c1) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %pb = fly.add_offset(%p0, %c2) : (!fly.ptr, !fly.int_tuple<128>) -> !fly.ptr + %a = fly.ptr.load(%pa) : (!fly.ptr) -> f32 + %b = fly.ptr.load(%pb) : (!fly.ptr) -> f32 + %s = arith.addf %a, %b : f32 + return %s : f32 +} + +// A multi-use inner add_offset is left alone: swapping it would keep the old op +// alive and add a second add_offset instead of replacing one. The form stays +// non-canonical but must still never fuse a runtime offset with a constant. +// CHECK-LABEL: @test_add_offset_multi_use_inner_kept +func.func @test_add_offset_multi_use_inner_kept(%ptr: !fly.ptr, %x: i32) -> f32 { + %c1 = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + // CHECK-NOT: arith.addi + // CHECK: %[[P1:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + // CHECK: %[[P2:.*]] = fly.add_offset(%[[P1]], %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: fly.ptr.load(%[[P1]]) + // CHECK: fly.ptr.load(%[[P2]]) + %p1 = fly.add_offset(%ptr, %c1) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p2 = fly.add_offset(%p1, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %a = fly.ptr.load(%p1) : (!fly.ptr) -> f32 + %b = fly.ptr.load(%p2) : (!fly.ptr) -> f32 + %s = arith.addf %a, %b : f32 + return %s : f32 +} + +// Zero and negative static offsets: 0 + (-32) merges to a single -32 tail. +// Guards the offsetDiv == 0 branch in AddOffsetOp type inference. +// CHECK-LABEL: @test_add_offset_zero_and_negative +func.func @test_add_offset_zero_and_negative(%ptr: !fly.ptr, %x: i32) -> f32 { + %d = fly.make_int_tuple(%x) : (i32) -> !fly.int_tuple + %z = fly.make_int_tuple() : () -> !fly.int_tuple<0> + %n = fly.make_int_tuple() : () -> !fly.int_tuple<-32> + // CHECK-NOT: arith.addi + // CHECK: %[[BASE:.*]] = fly.add_offset(%{{.*}}, %{{.*}}) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + // CHECK: %[[C:.*]] = fly.make_int_tuple() : () -> !fly.int_tuple<-32> + // CHECK: fly.add_offset(%[[BASE]], %[[C]]) : (!fly.ptr, !fly.int_tuple<-32>) -> !fly.ptr + %p1 = fly.add_offset(%ptr, %d) : (!fly.ptr, !fly.int_tuple) -> !fly.ptr + %p2 = fly.add_offset(%p1, %z) : (!fly.ptr, !fly.int_tuple<0>) -> !fly.ptr + %p3 = fly.add_offset(%p2, %n) : (!fly.ptr, !fly.int_tuple<-32>) -> !fly.ptr + %v = fly.ptr.load(%p3) : (!fly.ptr) -> f32 + return %v : f32 +} + +// A multi-use inner op still fuses when neither offset is dynamic: fusion +// rewrites only the outer op, so the inner one stays valid for its other +// users. Only the swap needs a single-use inner op. +// CHECK-LABEL: @test_add_offset_multi_use_inner_still_fuses +func.func @test_add_offset_multi_use_inner_still_fuses(%ptr: !fly.ptr) -> f32 { + %a = fly.make_int_tuple() : () -> !fly.int_tuple<64> + %b = fly.make_int_tuple() : () -> !fly.int_tuple<32> + // The outer offset is folded into a single static tail off the base pointer, + // even though %p1 is also loaded from directly. + // CHECK: %[[C96:.*]] = fly.make_int_tuple() : () -> !fly.int_tuple<96> + // CHECK: fly.add_offset(%{{.*}}, %[[C96]]) : (!fly.ptr, !fly.int_tuple<96>) -> !fly.ptr + %p1 = fly.add_offset(%ptr, %a) : (!fly.ptr, !fly.int_tuple<64>) -> !fly.ptr + %p2 = fly.add_offset(%p1, %b) : (!fly.ptr, !fly.int_tuple<32>) -> !fly.ptr + %x = fly.ptr.load(%p1) : (!fly.ptr) -> f32 + %y = fly.ptr.load(%p2) : (!fly.ptr) -> f32 + %s = arith.addf %x, %y : f32 + return %s : f32 +}