)]}'
{
  "commit": "af08a7c817dbfafbd6872670e2c792d4732319f7",
  "tree": "d4885d391a228601ed8e31a3c1fffc237a38e93a",
  "parents": [
    "be57650d3723e4b8b44a341bd007fcb02062dd5b"
  ],
  "author": {
    "name": "juanigp",
    "email": "pisula@roofline.ai",
    "time": "Wed Jul 01 17:16:26 2026 +0200"
  },
  "committer": {
    "name": "GitHub",
    "email": "noreply@github.com",
    "time": "Wed Jul 01 15:16:26 2026 +0000"
  },
  "message": "[DispatchCreation] Hoist scalar tensor.extract and tensor.extract_slice (#24552)\n\nThis PR aims to :\n- Enable hoisting tensor.extract ops that read from a scalar tensor\nalready outside the dispatch.\n- Loosen the requirement to hoist tensor.extract_slice ops.\n\n### Context\nI encountered this error when working on a modified version of LFM2.5. I\nattach a reproducer which captures the idea.\n\nA causal mask path stayed fused inside a QK matmul dispatch, which\nproduced large vectors due to tile size propagation in an mmt4d ukernel.\n\n```\nmask_slice_qk_repro.mlir:45:18: error: One or more operations with large vector sizes (32768 bytes) were found:\n\n\t%scores_3d \u003d torch.aten.bmm %q, %k : !torch.vtensor\u003c[16,?,64],f32\u003e, !torch.vtensor\u003c[16,64,256],f32\u003e -\u003e !torch.vtensor\u003c[16,?,256],f32\u003e\n                 ^\n\u003cunknown\u003e:0: note:   %cst \u003d arith.constant dense\u003c0xFF800000\u003e : vector\u003c16x16x16x16xf32\u003e\n\n\u003cunknown\u003e:0: note:   %cst_0 \u003d arith.constant dense\u003c0.000000e+00\u003e : vector\u003c16x16x16x16xf32\u003e\n\n\u003cunknown\u003e:0: note:   %cst_1 \u003d arith.constant dense\u003c0\u003e : vector\u003c16x16x16x16xi8\u003e\n```\n\nThis originated due to the mask\u0027s slicing offset coming from a\ntensor.extract and subsequently producing a scalar metadata chain that\nremained inside the dispatch.\n\n### Proposed fix\n- `HoistUniformScalarComputePass` can accept more \"candidate ops\" than\njust arith ops. I just included tensor.extract since its what I ran\ninto. `isUniformScalarForDispatch` is still in charge to verify that the\ncandidate op is hoistable, so i added the logic to check the\ntensor.extract ops.\n- `IREE::Flow::isOffsetSizeAndStrideMappableToFlow` got split into two:\n`isOffsetSizeAndStrideStructurallyMappableToFlow` just checks if the\nslice can be represented as one flat contiguous byte range, and\n`isOffsetSizeAndStrideMappableToFlow` checks for that and the additional\ntensor.extract provenance.\n- `isHoistableOp` in HoistEncodingOps.cpp was rejecting extract slice\nops whose offset, size, and stride where produced by an extract op due\nto calling `isOffsetSizeAndStrideMappableToFlow` on them. Now it calls\n`isOffsetSizeAndStrideStructurallyMappableToFlow`.\n\n### Additional Notes:\n- I tried to not interfere with the codebase\u0027s original intentions.\n- Since the extract and extract_slice make it out of the dispatch, the\nlarge vectors never occur. I thought this was the right way to address\nthe root cause of the problem.\n- Inspecting the mmt4d ukernel tile size propagation, it seems that the\nproblematic large vectors originated due to propagating a pack op tiling\nconfig to the outer dims of an accumulator, which should not happen\nafaiu. I could work on that separate issue if it is of interest.\n\n\u003cdetails\u003e\n\u003csummary\u003emlir reproducer\u003c/summary\u003e\n\n```mlir\nmodule @module {\n  func.func @forward(\n      %query: !torch.vtensor\u003c[1,16,?,64],f32\u003e,\n      %key: !torch.vtensor\u003c[1,16,64,256],f32\u003e,\n      %mask: !torch.vtensor\u003c[256,256],ui8\u003e,\n      %positions: !torch.vtensor\u003c[1],si64\u003e)\n      -\u003e !torch.vtensor\u003c[1,16,?,256],f32\u003e\n      attributes {torch.assume_strict_symbolic_shapes} {\n    %s \u003d torch.symbolic_int \"s\" {min_val \u003d 1, max_val \u003d 256} : !torch.int\n    torch.bind_symbolic_shape %query, [%s], affine_map\u003c()[s0] -\u003e (1, 16, s0, 64)\u003e : !torch.vtensor\u003c[1,16,?,64],f32\u003e\n\n    %int0 \u003d torch.constant.int 0\n    %int1 \u003d torch.constant.int 1\n    %int2 \u003d torch.constant.int 2\n    %int16 \u003d torch.constant.int 16\n    %int64 \u003d torch.constant.int 64\n    %int256 \u003d torch.constant.int 256\n    %int-1 \u003d torch.constant.int -1\n\n    %seq_len \u003d torch.aten.size.int %query, %int2 : !torch.vtensor\u003c[1,16,?,64],f32\u003e, !torch.int -\u003e !torch.int\n    %pos_tensor \u003d torch.aten.select.int %positions, %int0, %int-1 : !torch.vtensor\u003c[1],si64\u003e, !torch.int, !torch.int -\u003e !torch.vtensor\u003c[],si64\u003e\n    %pos \u003d torch.aten.item %pos_tensor : !torch.vtensor\u003c[],si64\u003e -\u003e !torch.int\n\n    %end \u003d torch.aten.add.int %pos, %seq_len : !torch.int, !torch.int -\u003e !torch.int\n    %bool_dtype \u003d torch.constant.int 11\n    %mask_bool \u003d torch.prims.convert_element_type %mask, %bool_dtype : !torch.vtensor\u003c[256,256],ui8\u003e, !torch.int -\u003e !torch.vtensor\u003c[256,256],i1\u003e\n    %mask_slice \u003d torch.aten.slice.Tensor %mask_bool, %int0, %pos, %end, %int1 : !torch.vtensor\u003c[256,256],i1\u003e, !torch.int, !torch.int, !torch.int, !torch.int -\u003e !torch.vtensor\u003c[?,256],i1\u003e\n    torch.bind_symbolic_shape %mask_slice, [%s], affine_map\u003c()[s0] -\u003e (s0, 256)\u003e : !torch.vtensor\u003c[?,256],i1\u003e\n\n    %float-Inf \u003d torch.constant.float 0xFFF0000000000000\n    %float0 \u003d torch.constant.float 0.000000e+00\n    %f32_dtype \u003d torch.constant.int 6\n    %none \u003d torch.constant.none\n    %cpu \u003d torch.constant.device \"cpu\"\n    %neg_inf \u003d torch.aten.scalar_tensor %float-Inf, %f32_dtype, %none, %cpu, %none : !torch.float, !torch.int, !torch.none, !torch.Device, !torch.none -\u003e !torch.vtensor\u003c[],f32\u003e\n    %mask_bias \u003d torch.aten.where.ScalarSelf %mask_slice, %float0, %neg_inf : !torch.vtensor\u003c[?,256],i1\u003e, !torch.float, !torch.vtensor\u003c[],f32\u003e -\u003e !torch.vtensor\u003c[?,256],f32\u003e\n    torch.bind_symbolic_shape %mask_bias, [%s], affine_map\u003c()[s0] -\u003e (s0, 256)\u003e : !torch.vtensor\u003c[?,256],f32\u003e\n\n    %q_shape \u003d torch.prim.ListConstruct %int16, %seq_len, %int64 : (!torch.int, !torch.int, !torch.int) -\u003e !torch.list\u003cint\u003e\n    %q \u003d torch.aten.view %query, %q_shape : !torch.vtensor\u003c[1,16,?,64],f32\u003e, !torch.list\u003cint\u003e -\u003e !torch.vtensor\u003c[16,?,64],f32\u003e\n    torch.bind_symbolic_shape %q, [%s], affine_map\u003c()[s0] -\u003e (16, s0, 64)\u003e : !torch.vtensor\u003c[16,?,64],f32\u003e\n\n    %k_shape \u003d torch.prim.ListConstruct %int16, %int64, %int256 : (!torch.int, !torch.int, !torch.int) -\u003e !torch.list\u003cint\u003e\n    %k \u003d torch.aten.view %key, %k_shape : !torch.vtensor\u003c[1,16,64,256],f32\u003e, !torch.list\u003cint\u003e -\u003e !torch.vtensor\u003c[16,64,256],f32\u003e\n    %scores_3d \u003d torch.aten.bmm %q, %k : !torch.vtensor\u003c[16,?,64],f32\u003e, !torch.vtensor\u003c[16,64,256],f32\u003e -\u003e !torch.vtensor\u003c[16,?,256],f32\u003e\n    torch.bind_symbolic_shape %scores_3d, [%s], affine_map\u003c()[s0] -\u003e (16, s0, 256)\u003e : !torch.vtensor\u003c[16,?,256],f32\u003e\n\n    %scores_shape \u003d torch.prim.ListConstruct %int1, %int16, %seq_len, %int256 : (!torch.int, !torch.int, !torch.int, !torch.int) -\u003e !torch.list\u003cint\u003e\n    %scores \u003d torch.aten.view %scores_3d, %scores_shape : !torch.vtensor\u003c[16,?,256],f32\u003e, !torch.list\u003cint\u003e -\u003e !torch.vtensor\u003c[1,16,?,256],f32\u003e\n    torch.bind_symbolic_shape %scores, [%s], affine_map\u003c()[s0] -\u003e (1, 16, s0, 256)\u003e : !torch.vtensor\u003c[1,16,?,256],f32\u003e\n\n    %result \u003d torch.aten.add.Tensor %scores, %mask_bias, %int1 : !torch.vtensor\u003c[1,16,?,256],f32\u003e, !torch.vtensor\u003c[?,256],f32\u003e, !torch.int -\u003e !torch.vtensor\u003c[1,16,?,256],f32\u003e\n    torch.bind_symbolic_shape %result, [%s], affine_map\u003c()[s0] -\u003e (1, 16, s0, 256)\u003e : !torch.vtensor\u003c[1,16,?,256],f32\u003e\n    return %result : !torch.vtensor\u003c[1,16,?,256],f32\u003e\n  }\n}\n\n```\n\u003c/details\u003e\n\nCompile command:\n```\niree-compile \\\n  mask_slice_qk_repro.mlir  \\\n  -o mask_slice_qk_repro.vmfb  \\\n  --iree-input-type\u003dauto \\\n  --iree-hal-target-device\u003dlocal \\\n  --iree-opt-data-tiling\u003dtrue \\\n  --iree-llvmcpu-enable-ukernels\u003dall \\\n  --iree-hal-local-target-device-backends\u003dllvm-cpu \\\n  --iree-hal-local-host-device-backends\u003dllvm-cpu \\\n  --iree-llvmcpu-target-cpu-features\u003dhost \n```\n\nAssisted by Codex 5.5\n\n---------\n\nSigned-off-by: Juan Ignacio Pisula \u003cpisula@roofline.ai\u003e",
  "tree_diff": [
    {
      "type": "modify",
      "old_id": "aef20cd47dc7bfc1018eef95eb1d8445c1e2044f",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp",
      "new_id": "5f8520b9745c4d38ec8327415ddeaa12a5cf294f",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.cpp"
    },
    {
      "type": "modify",
      "old_id": "071369809cb1ef868d12b594f696b29ad824f574",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h",
      "new_id": "592277e9b28a505beda76e17c1f5ea0f7277ac57",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h"
    },
    {
      "type": "modify",
      "old_id": "d225a32daec91c99f41395bd79eb91ab96f5a6f0",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp",
      "new_id": "47966028a2fbd7658bd753c4e6ab39cdbb44f036",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp"
    },
    {
      "type": "modify",
      "old_id": "e0c60964ef35a576f7e471943a6dbdfe390e8a02",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/DispatchCreation/HoistUniformScalarCompute.cpp",
      "new_id": "382627e861f7b773244f872347d873c6810dad02",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/DispatchCreation/HoistUniformScalarCompute.cpp"
    },
    {
      "type": "modify",
      "old_id": "2d075f5286ab70dfb18a92f4b6a52248fac86add",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir",
      "new_id": "95a1a7ac35b6b5da2816ba52e2846172a08a0a53",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/DispatchCreation/test/hoist_encoding_ops.mlir"
    },
    {
      "type": "modify",
      "old_id": "a29880f7aab0ec09b170e5ecdc310e0a442eb86c",
      "old_mode": 33188,
      "old_path": "compiler/src/iree/compiler/DispatchCreation/test/hoist_uniform_scalar_compute.mlir",
      "new_id": "365a1fdd25e48d6de89d258d4276f733465f824d",
      "new_mode": 33188,
      "new_path": "compiler/src/iree/compiler/DispatchCreation/test/hoist_uniform_scalar_compute.mlir"
    },
    {
      "type": "modify",
      "old_id": "8c0092137b7523cc40711d2f73eaa4950f1294c9",
      "old_mode": 33188,
      "old_path": "tests/e2e/stablehlo_ops/CMakeLists.txt",
      "new_id": "9ff2976cdc11fa13867597b4973a2593ccdfc2dc",
      "new_mode": 33188,
      "new_path": "tests/e2e/stablehlo_ops/CMakeLists.txt"
    }
  ]
}