[mlir][DialectConversion] Do not prematurely drop unused cast operations

[herhut]

Summary

The rewrite logic has an optimization to drop a cast operation after
rewriting block arguments if the cast operation has no users. This is
unsafe as there might be a pending rewrite that replaced the cast operation
itself and hence would trigger a second free.

Instead, do not remove the casts and leave it up to a later canonicalization
to do so.

[herhut]

PTAL. This fixes a use after free.

[bkramer]

Let's land this now to unblock asan. It's breaking test/Transforms/test-legalizer.mlir though because there's a cast being left behind.

[herhut]

Fix test

mlir/lib/Transforms/Utils/DialectConversion.cpp

for (unsigned i = 0, e = origBlock->getNumArguments(); i != e; ++i) {
 
       // Otherwise this is a 1->1+ value mapping.
       Value castValue = argInfo->castValue;
       assert(argInfo->newArgSize >= 1 && castValue && "expected 1->1+ mapping");
 
       // If the argument is still used, replace it with the generated cast.
       if (!origArg.use_empty())
         origArg.replaceAllUsesWith(mapping.lookupOrDefault(castValue));
-
-      // If all users of the cast were removed, we can drop it. Otherwise, keep
-      // the operation alive and let the user handle any remaining usages.
-      if (castValue.use_empty() && castValue.getDefiningOp())
-        castValue.getDefiningOp()->erase();
     }
   }
 }
 
 LogicalResult ArgConverter::materializeLiveConversions(
     ConversionValueMapping &mapping, OpBuilder &builder,
     function_ref<Operation *(Value)> findLiveUser) {
   for (auto &info : conversionInfo) {

mlir/test/Transforms/test-legalizer.mlir

func @remap_call_1_to_1(%arg0: i64) {
   // expected-remark@+1 {{op 'std.call' is not legalizable}}
   call @remap_input_1_to_1(%arg0) : (i64) -> ()
   // expected-remark@+1 {{op 'std.return' is not legalizable}}
   return
 }
 
 // CHECK-LABEL: func @remap_input_1_to_N({{.*}}f16, {{.*}}f16)
 func @remap_input_1_to_N(%arg0: f32) -> f32 {
+  // CHECK-NEXT: [[CAST:%.*]] = "test.cast"(%arg0, %arg1) : (f16, f16) -> f32
   // CHECK-NEXT: "test.return"{{.*}} : (f16, f16) -> ()
   "test.return"(%arg0) : (f32) -> ()
 }
 
 // CHECK-LABEL: func @remap_input_1_to_N_remaining_use(%arg0: f16, %arg1: f16)
 func @remap_input_1_to_N_remaining_use(%arg0: f32) {
   // CHECK-NEXT: [[CAST:%.*]] = "test.cast"(%arg0, %arg1) : (f16, f16) -> f32
   // CHECK-NEXT: "work"([[CAST]]) : (f32) -> ()
   // expected-remark@+1 {{op 'work' is not legalizable}}
   "work"(%arg0) : (f32) -> ()