EnzymeAD / EnzymeAD/Enzyme-JAX

Gather/scatter emitted by gb main ci

Open
#2,522 10 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

```

%c_108 = stablehlo.constant dense<16> : tensor<256xi64> loc(#loc645)

%c_110 = stablehlo.constant dense<"0xFFFFFFFFFFFFFFFF0100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F0000000000000080000000000000008100000000000000820000000000000083000000000000008400000000000000850000000000000086000000000000008700000000000000880000000000000089000000000000008A000000000000008B000000000000008C000000000000008D000000000000008E000000000000008F0000000000000090000000000000009100000000000000920000000000000093000000000000009400000000000000950000000000000096000000000000009700000000000000980000000000000099000000000000009A000000000000009B000000000000009C000000000000009D000000000000009E000000000000009F00000000000000A000000000000000A100000000000000A200000000000000A300000000000000A400000000000000A500000000000000A600000000000000A700000000000000A800000000000000A900000000000000AA00000000000000AB00000000000000AC00000000000000AD00000000000000AE00000000000000AF00000000000000B000000000000000B100000000000000B200000000000000B300000000000000B400000000000000B500000000000000B600000000000000B700000000000000B800000000000000B900000000000000BA00000000000000BB00000000000000BC00000000000000BD00000000000000BE00000000000000BF00000000000000C000000000000000C100000000000000C200000000000000C300000000000000C400000000000000C500000000000000C600000000000000C700000000000000C800000000000000C900000000000000CA00000000000000CB00000000000000CC00000000000000CD00000000000000CE00000000000000CF00000000000000D000000000000000D100000000000000D200000000000000D300000000000000D400000000000000D500000000000000D600000000000000D700000000000000D800000000000000D900000000000000DA00000000000000DB00000000000000DC00000000000000DD00000000000000DE00000000000000DF00000000000000E000000000000000E100000000000000E200000000000000E300000000000000E400000000000000E500000000000000E600000000000000E700000000000000E800000000000000E900000000000000EA00000000000000EB00000000000000EC00000000000000ED00000000000000EE00000000000000EF00000000000000F000000000000000F100000000000000F200000000000000F300000000000000F400000000000000F500000000000000F600000000000000F700000000000000F800000000000000F900000000000000FA00000000000000FB00000000000000FC00000000000000FD00000000000000FE00000000000000FF00000000000000"> : tensor<256xi64> loc(#loc645)

%c_88 = stablehlo.constant {enzymexla.non_negative = [#enzymexla]} dense<1> : tensor<256xi64> loc(#loc)

%c_109 = stablehlo.constant dense<"0x01000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"> : tensor<256xi1> loc(#loc645)

%c_85 = stablehlo.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94]> : tensor<95xi64> loc(#loc1135)

%c_332 = stablehlo.constant dense<16> : tensor<95xi64> loc(#loc925)

%25 = stablehlo.divide %c_110, %c_108 : tensor<256xi64> loc(#loc645)
%26 = stablehlo.negate %25 : tensor<256xi64> loc(#loc645)

%27 = stablehlo.subtract %26, %c_88 {enzymexla.non_negative = [#enzymexla]} : tensor<256xi64> loc(#loc645)

%28 = stablehlo.select %c_109, %27, %25 {enzymexla.non_negative = [#enzymexla]} : tensor<256xi1>, tensor<256xi64> loc(#loc645)

%3067 = stablehlo.multiply %c_85, %c_332 : tensor<95xi64> loc(#loc1162)

%3068 = stablehlo.broadcast_in_dim %3067, dims = [0] : (tensor<95xi64>) -> tensor<95x256xi64> loc(#loc925)
%3069 = stablehlo.broadcast_in_dim %28, dims = [1] : (tensor<256xi64>) -> tensor<95x256xi64> loc(#loc925)

%3070 = stablehlo.add %3068, %3069 : tensor<95x256xi64> loc(#loc925)
%3071 = stablehlo.reshape %3070 : (tensor<95x256xi64>) -> tensor<24320x1xi64> loc(#loc925)

%3083 = stablehlo.add %3071, %c_334 : tensor<24320x1xi64> loc(#loc925)
%3084 = "stablehlo.gather"(%arg6, %3083) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<1536xf32>, tensor<24320x1xi64>) -> tensor<24320xf32> loc(#loc925)
```

cc @avik-pal @dkytezab . Bad regardless, may be part of the cause of current pain

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the StableHLO snippet in this issue and reproduce the gb main CI case that emits it. Trace the generation of the gather and preceding index computations, then compare the emitted IR with the intended gather/scatter behavior. Done means the problematic emission is explained and corrected or a focused reproducer and actionable diagnosis is added.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.