[PATCH] D87746: [mlir] Fix of updating function signature in normalizing memrefs

[Haruki Imai via Phabricator via llvm-commits]

imaihal updated this revision to Diff 292512.
imaihal added a comment.

Updated comments and rename test func


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D87746/new/

https://reviews.llvm.org/D87746

Files:
  mlir/lib/Transforms/NormalizeMemRefs.cpp
  mlir/test/Transforms/normalize-memrefs-ops.mlir
  mlir/test/lib/Dialect/Test/TestDialect.h
  mlir/test/lib/Dialect/Test/TestOps.td


Index: mlir/test/lib/Dialect/Test/TestOps.td
===================================================================
--- mlir/test/lib/Dialect/Test/TestOps.td
+++ mlir/test/lib/Dialect/Test/TestOps.td
@@ -628,6 +628,17 @@
   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.
Index: mlir/test/lib/Dialect/Test/TestDialect.h
===================================================================
--- mlir/test/lib/Dialect/Test/TestDialect.h
+++ mlir/test/lib/Dialect/Test/TestDialect.h
@@ -16,6 +16,7 @@
 
 #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"
@@ -29,7 +30,6 @@
 
 #include "TestOpEnums.h.inc"
 
-
 #include "TestOpStructs.h.inc"
 #include "TestOpsDialect.h.inc"
 
Index: mlir/test/Transforms/normalize-memrefs-ops.mlir
===================================================================
--- mlir/test/Transforms/normalize-memrefs-ops.mlir
+++ mlir/test/Transforms/normalize-memrefs-ops.mlir
@@ -55,3 +55,6 @@
     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
+
+// Test with an arbitrary op that references the function symbol.
+"test.op_funcref"() {func = @test_norm_mix} : () -> ()
Index: mlir/lib/Transforms/NormalizeMemRefs.cpp
===================================================================
--- mlir/lib/Transforms/NormalizeMemRefs.cpp
+++ mlir/lib/Transforms/NormalizeMemRefs.cpp
@@ -263,16 +263,23 @@
   // 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
@@ -302,8 +309,8 @@
     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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D87746.292512.patch
Type: text/x-patch
Size: 4131 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200917/f044e7f8/attachment.bin>