[MLIR] Fix AsmPrinter.cpp for single ssa-id AffineMap

The AsmPrinter wrongly assumes that all single ssa-id AffineMap
are the identity map for the purpose of printing.
This CL adds the missing level of indirection as well as a test.

This bug was originally shaken off by the experimental TC->MLIR path.

Before this CL, the test would print:
```
mlfunc @mlfuncsimplemap(%arg0 : affineint, %arg1 : affineint, %arg2 : affineint) {
  for %i0 = 0 to %arg0 {
    for %i1 = 0 to %i0 {
                   ~~~   should be %arg1
      %c42_i32 = constant 42 : i32
    }
  }
  return
}
```

PiperOrigin-RevId: 214120817
This commit is contained in:
Nicolas Vasilache 2018-09-22 10:35:44 -07:00 committed by jpienaar
parent e5354c2404
commit f9e50199e9
2 changed files with 20 additions and 3 deletions

View File

@ -1431,9 +1431,14 @@ void MLFunctionPrinter::printBound(AffineBound bound, const char *prefix) {
return;
}
// Print bound that consists of a single SSA id.
if (isa<AffineDimExpr>(expr) || isa<AffineSymbolExpr>(expr)) {
printOperand(bound.getOperand(0));
// Print bound that consists of a single SSA id, we need an indirection
// to achieve this.
if (auto *dimExpr = dyn_cast<AffineDimExpr>(expr)) {
printOperand(bound.getOperand(dimExpr->getPosition()));
return;
} else if (auto *symExpr = dyn_cast<AffineSymbolExpr>(expr)) {
printOperand(
bound.getOperand(map->getNumDims() + symExpr->getPosition()));
return;
}
} else {

View File

@ -454,3 +454,15 @@ mlfunc @mlfuncattrempty() -> ()
attributes {} {
return
}
// CHECK-label mlfunc @mlfuncsimplemap
#mapsimple0 = ()[s0, s1, s2] -> (s0)
#mapsimple1 = (d0)[s0, s1, s2] -> (s1)
mlfunc @mlfuncsimplemap(%arg0 : affineint, %arg1 : affineint, %arg2 : affineint) -> () {
for %i0 = 0 to #mapsimple0()[%arg0, %arg1, %arg2] { // CHECK: for %i0 = 0 to %arg0 {
for %i1 = 0 to #mapsimple1(%i0)[%arg0, %arg1, %arg2] { // CHECK: for %i1 = 0 to %arg1 {
%c42_i32 = constant 42 : i32
}
}
return
}