Skip to content

linalg.index shouldn't be marked as Pure. #62644

@erick-xanadu

Description

@erick-xanadu

Currently linalg.index is marked as Pure which leads to different linalg.index operations to be eliminated via CSE. This will mean that linalg.index operations with the same operands yet on different enclosing scopes will be transformed leading to bad iteration indices.

Consider the following example:

 linalg.map outs(%alloc_4 : memref<2x3x2xf64>) () { %3 = linalg.index 0 : index %4 = linalg.index 1 : index %5 = linalg.index 2 : index // ... snip ...  linalg.generic {indexing_maps = [#map1], iterator_types = ["parallel", "parallel"]} outs(%alloc_12 : memref<3x2xf64>) { ^bb0(%out: f64): %9 = linalg.index 1 : index %10 = linalg.index 0 : index %11 = arith.cmpi ult, %10, %c1 : index %12 = scf.if %11 -> (f64) { %13 = memref.load %expand_shape_9[%10, %9] : memref<1x2xf64>  scf.yield %13 : f64 } else {

After applying CSE linalg.index ops inside linalg.generic get eliminated:

 linalg.map outs(%alloc_4 : memref<2x3x2xf64>) () { %3 = linalg.index 0 : index %4 = linalg.index 1 : index %5 = linalg.index 2 : index // ... snip ...  linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} outs(%alloc_12 : memref<3x2xf64>) { ^bb0(%out: f64): %9 = arith.cmpi ult, %3, %c1 : index %10 = scf.if %9 -> (f64) { %11 = memref.load %expand_shape_9[%3, %4] : memref<1x2xf64>  scf.yield %11 : f64

Later after lowering to loops:

 scf.for %arg1 = %c0 to %c2 step %c1 { scf.for %arg2 = %c0 to %c3 step %c1 { scf.for %arg3 = %c0 to %c2 step %c1 { // ...snip... scf.for %arg4 = %c0 to %c3 step %c1 { scf.for %arg5 = %c0 to %c2 step %c1 { %6 = arith.cmpi ult, %arg1, %c1 : index %7 = scf.if %6 -> (f64) { %8 = memref.load %expand_shape_9[%arg1, %arg2] : memref<1x2xf64> 

But the last line should be memref.load %expanded_shape_9[%arg4, %arg5]

Metadata

Metadata

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions