Do not use extractvalue if the inserted value is directly reachable

[AndreyPavlenko]

No description provided.

[alexbaden]

Can you provide some description here to include the why, not just the what or the how.

[AndreyPavlenko]

Can you provide some description here to include the why, not just the what or the how.

Sure, I'll add when it's ready for review. It's a draft and will definitely be changed (or closed, because I'm not quite sure about it).

[alexbaden]

I'm not quite sure about it either :) some more context w/ the draft might help generate ideas for how to proceed.

[AndreyPavlenko]

some more context w/ the draft might help generate ideas for how to proceed.

This is an attempt to minimise the number of insert/extract_value operations - #4136.

Here #4062 (comment) is an example with a huge amount of such operations, where constants are inserted and then extracted from structures. Later all these operations are optimised out by the canonicalizer, but before that a huger IR is created.

[AndreyPavlenko]

A short example:

#blocked0 = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0], CTAsPerCGA = [1], CTASplitNum = [1], CTAOrder = [0]}>
#blocked2 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [0, 1], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [0, 1]}>
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
  tt.func @basic_view_broadcast(%arg : tensor<256xf32,#blocked0>) {
    %0 = tt.reshape %arg allow_reorder : tensor<256xf32, #blocked0> -> tensor<256x1xf32,#blocked2>
    %1 = tt.broadcast %0 : tensor<256x1xf32,#blocked2> -> tensor<256x4xf32, #blocked2>
    tt.return
  }
}

lowered to:

module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
  llvm.mlir.global external @global_smem() {addr_space = 3 : i32, alignment = 16 : i64} : !llvm.array<0 x i8>
  llvm.func spir_kernelcc @basic_view_broadcast(%arg0: !llvm.struct<(f32, f32)>) attributes {intel_reqd_sub_group_size = 32 : i32, triton_gen.max_work_group_size = array<i32: 128, 1, 1>} {
    %0 = llvm.extractvalue %arg0[0] : !llvm.struct<(f32, f32)>
    %1 = llvm.extractvalue %arg0[1] : !llvm.struct<(f32, f32)>
    %2 = llvm.mlir.undef : !llvm.struct<(f32, f32)>
    %3 = llvm.insertvalue %0, %2[0] : !llvm.struct<(f32, f32)>
    %4 = llvm.insertvalue %1, %3[1] : !llvm.struct<(f32, f32)>
    %5 = llvm.extractvalue %4[0] : !llvm.struct<(f32, f32)>
    %6 = llvm.extractvalue %4[1] : !llvm.struct<(f32, f32)>
    %7 = llvm.mlir.undef : !llvm.struct<(f32, f32, f32, f32)>
    %8 = llvm.insertvalue %5, %7[0] : !llvm.struct<(f32, f32, f32, f32)>
    %9 = llvm.insertvalue %6, %8[1] : !llvm.struct<(f32, f32, f32, f32)>
    %10 = llvm.insertvalue %5, %9[2] : !llvm.struct<(f32, f32, f32, f32)>
    %11 = llvm.insertvalue %6, %10[3] : !llvm.struct<(f32, f32, f32, f32)>
    llvm.return
  }
}

The lines

    %3 = llvm.insertvalue %0, %2[0] : !llvm.struct<(f32, f32)>
    %4 = llvm.insertvalue %1, %3[1] : !llvm.struct<(f32, f32)>
    %5 = llvm.extractvalue %4[0] : !llvm.struct<(f32, f32)>
    %6 = llvm.extractvalue %4[1] : !llvm.struct<(f32, f32)>

are redundant. This fix checks if the structure is llvm.insertvalue, it returns the inserted value, instead of creating llvm.extractvalue operation.

[Copilot]

Pull Request Overview

This PR refactors the unpacking function in Utility.cpp to avoid using extract_val on elements that are directly available from the InsertValueOp chain.

• Introduces a loop to traverse the chain of InsertValueOps and assign inserted values directly.
• Falls back to extract_val for any remaining elements.

Files not reviewed (7)

• test/Conversion/amd/buffer_load_store.mlir: Language not supported
• test/Conversion/intel/dot_layout_offset.mlir: Language not supported
• test/Conversion/intel/tritongpu_to_gen.mlir: Language not supported
• test/Conversion/tritongpu_to_llvm.mlir: Language not supported
• test/TritonIntelGPU/blockptr_load.mlir: Language not supported
• test/TritonIntelGPU/blockptr_store.mlir: Language not supported
• test/TritonIntelGPU/prefetch-to-llvm.mlir: Language not supported

Comments suppressed due to low confidence (1)

lib/Conversion/TritonGPUToLLVM/Utility.cpp:567

• [nitpick] Enhance this comment to clarify the assumption that the InsertValueOp chain is contiguous, ensuring future maintainers understand the logic underlying the loop condition.

// If llvmStruct is an InsertValueOp, iterate up over the chain of InsertValueOps and get the inserted values instead of extracting from the struct.

[Copilot]

Ensure that types is non-empty before this loop to avoid unsigned underflow in 'remaining - 1'. Consider adding an early-return or guard when types.size() is zero.