[MLIR] normalize-memrefs: Ilegal IR on memref.load/memref.store

[mgehre-amd]

-normalize-memrefs on

#map = affine_map<(d0, d1) -> (d1, d0)>
func.func @forward(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref<85x40xf32, #map>
    %8 = memref.load %alloc[%arg0, %arg1] : memref<85x40xf32, #map>
    return
}

produces

  func.func @forward(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref<40x85xf32>
    %0 = memref.load %alloc[%arg0, %arg1] : memref<40x85xf32>
    return
  }

where the indices into %alloc within the memref.load have not been swapped as they should be.
I would expect the result to be

  func.func @forward(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref<40x85xf32>
    %0 = memref.load %alloc[%arg1, %arg0] : memref<40x85xf32>
    return
  }

As comparison, applying this to the same IR with affine.load:

func.func @affine(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref<85x40xf32, #map>
    %8 = affine.load %alloc[%arg0, %arg1] : memref<85x40xf32, #map>
    return
}

results in a correct swap of the indices:

  func.func @affine(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref<40x85xf32>
    %0 = affine.load %alloc[symbol(%arg1), symbol(%arg0)] : memref<40x85xf32>
    return
  }

Reproducer: https://godbolt.org/z/5Pffq5fKP

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Matthias Gehre (mgehre-amd)

`-normalize-memrefs` on ``` #map = affine_map<(d0, d1) -> (d1, d0)> func.func @forward(%arg0: index, %arg1: index) { %alloc = memref.alloc() : memref<85x40xf32, #map> %8 = memref.load %alloc[%arg0, %arg1] : memref<85x40xf32, #map> return } ``` produces ``` func.func @forward(%arg0: index, %arg1: index) { %alloc = memref.alloc() : memref<40x85xf32> %0 = memref.load %alloc[%arg0, %arg1] : memref<40x85xf32> return } ``` where the indices into `%alloc` within the memref.load have not been swapped as they should be. I would expect the result to be ``` func.func @forward(%arg0: index, %arg1: index) { %alloc = memref.alloc() : memref<40x85xf32> %0 = memref.load %alloc[%arg1, %arg0] : memref<40x85xf32> return } ```

As comparison, applying this to the same IR with affine.load:

func.func @<!-- -->affine(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref&lt;85x40xf32, #map&gt;
    %8 = affine.load %alloc[%arg0, %arg1] : memref&lt;85x40xf32, #map&gt;
    return
}

results in a correct swap of the indices:

  func.func @<!-- -->affine(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref&lt;40x85xf32&gt;
    %0 = affine.load %alloc[symbol(%arg1), symbol(%arg0)] : memref&lt;40x85xf32&gt;
    return
  }

[llvmbot]

@llvm/issue-subscribers-bug

Author: Matthias Gehre (mgehre-amd)

`-normalize-memrefs` on ``` #map = affine_map<(d0, d1) -> (d1, d0)> func.func @forward(%arg0: index, %arg1: index) { %alloc = memref.alloc() : memref<85x40xf32, #map> %8 = memref.load %alloc[%arg0, %arg1] : memref<85x40xf32, #map> return } ``` produces ``` func.func @forward(%arg0: index, %arg1: index) { %alloc = memref.alloc() : memref<40x85xf32> %0 = memref.load %alloc[%arg0, %arg1] : memref<40x85xf32> return } ``` where the indices into `%alloc` within the memref.load have not been swapped as they should be. I would expect the result to be ``` func.func @forward(%arg0: index, %arg1: index) { %alloc = memref.alloc() : memref<40x85xf32> %0 = memref.load %alloc[%arg1, %arg0] : memref<40x85xf32> return } ```

As comparison, applying this to the same IR with affine.load:

func.func @<!-- -->affine(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref&lt;85x40xf32, #map&gt;
    %8 = affine.load %alloc[%arg0, %arg1] : memref&lt;85x40xf32, #map&gt;
    return
}

results in a correct swap of the indices:

  func.func @<!-- -->affine(%arg0: index, %arg1: index) {
    %alloc = memref.alloc() : memref&lt;40x85xf32&gt;
    %0 = affine.load %alloc[symbol(%arg1), symbol(%arg0)] : memref&lt;40x85xf32&gt;
    return
  }

[CoTinker]

The code is here, for that memref.load not implement AffineMapAccessInterface.

llvm-project/mlir/lib/Dialect/Affine/Utils/Utils.cpp

Lines 1226 to 1247 in 36a0d44

	auto affMapAccInterface = dyn_cast<AffineMapAccessInterface>(op);
	if (!affMapAccInterface) {
	if (!allowNonDereferencingOps) {
	// Failure: memref used in a non-dereferencing context (potentially
	// escapes); no replacement in these cases unless allowNonDereferencingOps
	// is set.
	return failure();
	}
	// Check if it is a memref.load
	auto memrefLoad = dyn_cast<memref::LoadOp>(op);
	bool isReductionLike =
	indexRemap.getNumResults() < indexRemap.getNumInputs();
	if (!memrefLoad || !isReductionLike) {
	op->setOperand(memRefOperandPos, newMemRef);
	return success();
	}
	return transformMemRefLoadWithReducedRank(
	op, oldMemRef, newMemRef, memRefOperandPos, extraIndices, extraOperands,
	symbolOperands, indexRemap);
	}

[mgehre-amd]

I guess memref.load cannot implement AffineMapAccessInterface because it doesn't have an attribute that would carry an affine map for the access, right?

[CoTinker]

Yes, but we could consider resetting the indices of memref.load based on the map.

[mgehre-amd]

@CoTinker, if I understand this correctly, the normalize-memref pass would then need to have special cases for memref.load/store, correct?

[CoTinker]

Looks like it is.