[mlir] Fix of updating function signature in normalizing memrefs

[imaihal]

Summary

Normalizing memrefs failed when a caller of symbolic use in a function
can not be casted to CallOp. This patch avoids the failure by checking
the result of the casting. If the caller can not be casted to CallOp,
it is skipped.

[imaihal]

This bug was discussed here https://llvm.discourse.group/t/memrefs-and-maps-for-tiling/1279/21
Following example failed because spv.EntryPoint "GLCompute" @empty can't be casted to CallOp. I created the example to reproduce errors happened in my code which uses EntryPoint op of my own dialect. This example may not be practical, but this reproduced the same error with my code.

#map0 = affine_map<(d0, d1) -> (d0 floordiv 32, d1 floordiv 64, d0 mod 32, d1 mod 64)>

module {
  func @empty() {
    %0 = alloc() : memref<10x10xf32, #map0>
    return
  }
  spv.EntryPoint "GLCompute" @empty
}

@bondhugula @avarmapml Could you review this patch?

[bondhugula]

Thanks for raising this. Please add a test case. It's not fully clear one could just ignore all uses that aren't CallOps. Please also add a comment.

[bondhugula]

mlir/lib/Transforms/NormalizeMemRefs.cpp

266–276

Please rename variables for better readability.

callOp -> userOp
castCallOp -> callOp

Also, please invert the logic for clarity.

auto callOp = dyn_cast<CallOp>(userOp);
if (!callOp)
  continue;
StringRef callee = castCallOp.getCallee();

[imaihal]

Added test and comments, and renamee variables to reflrect comments

[avarmapml]

Thank you for raising this.
Solution looks good to me. Please add the test case though.

[imaihal]

@bondhugula @avarmapml Thanks for your comments! I added test and comments, and renamed the variables to reflect @bondhugula 's comments.

[avarmapml]

mlir/lib/Transforms/NormalizeMemRefs.cpp

271–272

You can rephrase this line as :-
TODO: Handle cases where a non-CallOp symbol use of a function deals with memrefs.

[bondhugula]

mlir/test/Transforms/normalize-memrefs-ops.mlir

59

-> Test with an arbitrary op that references the symbol
?

Nit: Terminate all comments with period.

mlir/test/lib/Dialect/Test/TestOps.td

632

Nit: terminate comment with a period.

632

including entrypoint -> including a reference to a function symbol.

643

Nit: Drop extra blank line.

[imaihal]

Rebased and rephrased the comment.

[imaihal]

Sorry, I may did wrong operaton... I will fix this.

[imaihal]

Rebased

[imaihal]

Fixed comments and format

[imaihal]

[bondhugula]

mlir/test/lib/Dialect/Test/TestOps.td

636

It's still not clear what this means. Did you mean to say it's an arbitrary op with a reference to a function symbol? Why is even an entry point in the context here?

[imaihal]

mlir/test/lib/Dialect/Test/TestOps.td

636

A op for entry point issued the error in my actual use case. So, I created the test op for entry point, but as you said, we should make it more general.

OPEntryPoint -> OPFuncRef
op_entrypoint -> op_funcref
The "test.op_entrypoint" function indicates the main entry point.
-> The "test.op_funcref" function has a reference to a function symbol
How about this?

[imaihal]

Updated comments and rename test func

[bondhugula]

Thanks!

[AlexEichenberger]

@imaihal Thanks for the upgrade. Minor nit still remains, please fix.

mlir/test/lib/Dialect/Test/TestOps.td

632

@imaihal Please add the period.

[imaihal]

mlir/test/lib/Dialect/Test/TestOps.td

632

@AlexEichenberger This is a two line-comment. Is it ok as it is?

// Test for memrefs normalization of an op with a reference to a function
// symbol.

Thanks!

[imaihal]

@bondhugula @AlexEichenberger @avarmapml Do you have any other comments?

I think I don't have right to commit. So, could you commit this patch if this is ready?

[imaihal]

@bondhugula It seems there is no additional comments. Could you commit also this patch?

[imaihal]

Rebased

mlir/lib/Transforms/NormalizeMemRefs.cpp

void NormalizeMemRefs::updateFunctionSignature(FuncOp funcOp,
   // function in ReturnOps, the caller function's signature will also change.
   // Hence we record the caller function in 'funcOpsToUpdate' to update their
   // signature as well.
   llvm::SmallDenseSet<FuncOp, 8> funcOpsToUpdate;
   // We iterate over all symbolic uses of the function and update the return
   // type at the caller site.
   Optional<SymbolTable::UseRange> symbolUses = funcOp.getSymbolUses(moduleOp);
   for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
-    Operation *callOp = symbolUse.getUser();
-    OpBuilder builder(callOp);
-    StringRef callee = cast<CallOp>(callOp).getCallee();
+    Operation *userOp = symbolUse.getUser();
+    OpBuilder builder(userOp);
+    // When `userOp` can not be casted to `CallOp`, it is skipped. This assumes
+    // that the non-CallOp has no memrefs to be replaced.
+    // TODO: Handle cases where a non-CallOp symbol use of a function deals with
+    // memrefs.
+    auto callOp = dyn_cast<CallOp>(userOp);
+    if (!callOp)
+      continue;
+    StringRef callee = callOp.getCallee();
     Operation *newCallOp = builder.create<CallOp>(
-        callOp->getLoc(), resultTypes, builder.getSymbolRefAttr(callee),
-        callOp->getOperands());
+        userOp->getLoc(), resultTypes, builder.getSymbolRefAttr(callee),
+        userOp->getOperands());
     bool replacingMemRefUsesFailed = false;
     bool returnTypeChanged = false;
-    for (unsigned resIndex : llvm::seq<unsigned>(0, callOp->getNumResults())) {
-      OpResult oldResult = callOp->getResult(resIndex);
+    for (unsigned resIndex : llvm::seq<unsigned>(0, userOp->getNumResults())) {
+      OpResult oldResult = userOp->getResult(resIndex);
       OpResult newResult = newCallOp->getResult(resIndex);
       // This condition ensures that if the result is not of type memref or if
       // the resulting memref was already having a trivial map layout then we
       // need not perform any use replacement here.
       if (oldResult.getType() == newResult.getType())
         continue;
       AffineMap layoutMap =
           oldResult.getType().dyn_cast<MemRefType>().getAffineMaps().front();
for (unsigned resIndex : llvm::seq<unsigned>(0, userOp->getNumResults())) {
         replacingMemRefUsesFailed = true;
         break;
       }
       returnTypeChanged = true;
     }
     if (replacingMemRefUsesFailed)
       continue;
     // Replace all uses for other non-memref result types.
-    callOp->replaceAllUsesWith(newCallOp);
-    callOp->erase();
+    userOp->replaceAllUsesWith(newCallOp);
+    userOp->erase();
     if (returnTypeChanged) {
       // Since the return type changed it might lead to a change in function's
       // signature.
       // TODO: If funcOp doesn't return any memref type then no need to update
       // signature.
       // TODO: Further optimization - Check if the memref is indeed part of
       // ReturnOp at the parentFuncOp and only then updation of signature is
       // required.

mlir/test/Transforms/normalize-memrefs-ops.mlir

func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
     dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
     // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x64xf32>
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
 
 // Test with maps in load and store ops.
 
 #map_tile = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 32, d2 mod 32, d3 mod 32)>
 
 // CHECK-LABEL: test_load_store
 // CHECK-SAME: (%[[ARG0:[a-z0-9]*]]: memref<1x16x14x14xf32>
 func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
     %0 = alloc() : memref<1x16x14x14xf32, #map_tile>
     // CHECK: %[[v0:[a-z0-9]*]] = alloc() : memref<1x16x1x1x32x32xf32>
affine.for %i = 0 to 1 {
       }
     }
     dealloc %0 :  memref<1x16x14x14xf32, #map_tile>
     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
     dealloc %1 :  memref<1x16x14x14xf32>
     // CHECK: dealloc %[[v1]] : memref<1x16x14x14xf32>
     return
 }
+
+// Test with an arbitrary op that references the function symbol.
+
+"test.op_funcref"() {func = @test_norm_mix} : () -> ()

mlir/test/lib/Dialect/Test/TestDialect.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTDIALECT_H
 #define MLIR_TESTDIALECT_H
 
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/Function.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/DerivedAttributeOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "TestOpEnums.h.inc"
 
-
 #include "TestOpStructs.h.inc"
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
 namespace mlir {
 void registerTestDialect(DialectRegistry &registry);
 } // end namespace mlir
 
 #endif // MLIR_TESTDIALECT_H

mlir/test/lib/Dialect/Test/TestOps.td

 def OpNorm : TEST_Op<"op_norm", [MemRefsNormalizable]> {
   let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
 }
 // Test for memrefs normalization of an op without normalizable memrefs.
 def OpNonNorm : TEST_Op<"op_nonnorm"> {
   let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
 }
 
+// Test for memrefs normalization of an op with a reference to a function
+// symbol.
+def OpFuncRef : TEST_Op<"op_funcref"> {
+  let summary = "Test op with a reference to a function symbol";
+  let description = [{
+    The "test.op_funcref" is a test op with a reference to a function symbol.
+  }];
+  let builders = [OpBuilder<[{OpBuilder &builder, OperationState &state,
+                             FuncOp function}]>];
+}
+
 // Pattern add the argument plus a increasing static number hidden in
 // OpMTest function. That value is set into the optional argument.
 // That way, we will know if operations is called once or twice.
 def OpMGetNullAttr : NativeCodeCall<"Attribute()">;
 def OpMAttributeIsNull : Constraint<CPred<"! ($_self)">, "Attribute is null">;
 def OpMVal : NativeCodeCall<"OpMTest($_builder, $0)">;
 def : Pat<(OpM $attr, $optAttr), (OpM $attr, (OpMVal $attr) ),
     [(OpMAttributeIsNull:$optAttr)]>;