Allow for remapping argument to a Value in SignatureConversion.

The current SignatureConversion framework (part of DialectConversion)
allows remapping input arguments to a function from 1->0, 1->1 or
1->many arguments during conversion. Another case is where the
argument itself is dropped, but it's use are remapped to another
Value*.

An example of this is: The Vulkan/SPIR-V spec requires entry functions
to be of type void(void). The GPU -> SPIR-V conversion implemented
this without having the DialectConversion framework track the
remapping that lead to some undefined behavior. The changes here
addresses that.

PiperOrigin-RevId: 275059656

mlir/g3doc/DialectConversion.md

@@ -222,6 +222,10 @@ public:
     /// new signature.
     void remapInput(unsigned origInputNo, unsigned newInputNo,
                     unsigned newInputCount = 1);
+
+    /// Remap an input of the original signature to another `replacement`
+    /// value. This drops the original argument.
+    void remapInput(unsigned origInputNo, Value *replacement);
 };
 ```
 

mlir/include/mlir/Transforms/DialectConversion.h

@@ -55,17 +55,18 @@ public:
     SignatureConversion(unsigned numOrigInputs)
         : remappedInputs(numOrigInputs) {}
 
-    /// This struct represents a range of new types that remap an existing
-    /// signature input.
+    /// This struct represents a range of new types or a single value that
+    /// remaps an existing signature input.
     struct InputMapping {
       size_t inputNo, size;
+      Value *replacementValue;
     };
 
     /// Return the argument types for the new signature.
     ArrayRef<Type> getConvertedTypes() const { return argTypes; }
 
     /// Get the input mapping for the given argument.
-    llvm::Optional<InputMapping> getInputMapping(unsigned input) const {
+    llvm::Optional<InputMapping> const &getInputMapping(unsigned input) const {
       return remappedInputs[input];
     }
 
@@ -86,6 +87,10 @@ public:
     void remapInput(unsigned origInputNo, unsigned newInputNo,
                     unsigned newInputCount = 1);
 
+    /// Remap an input of the original signature to another `replacement`
+    /// value. This drops the original argument.
+    void remapInput(unsigned origInputNo, Value *replacement);
+
   private:
     /// The remapping information for each of the original arguments.
     SmallVector<llvm::Optional<InputMapping>, 4> remappedInputs;

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

@@ -244,7 +244,7 @@ LogicalResult lowerAsEntryFunction(FuncOp funcOp,
     for (auto origArg : enumerate(funcOp.getArguments())) {
       auto replacement = createAndLoadGlobalVarForEntryFnArg(
           rewriter, origArg.index(), origArg.value());
-      rewriter.replaceUsesOfBlockArgument(origArg.value(), replacement);
+      signatureConverter.remapInput(origArg.index(), replacement);
     }
   }
   newFuncOp = applySignatureConversion(funcOp, rewriter, signatureConverter);

mlir/lib/Transforms/DialectConversion.cpp

@@ -288,8 +288,17 @@ void ArgConverter::applySignatureConversion(
   rewriter.setInsertionPointToStart(block);
   for (unsigned i = 0; i != origArgCount; ++i) {
     ArrayRef<Value *> remappedValues;
-    if (auto inputMap = signatureConversion.getInputMapping(i))
+    if (auto &inputMap = signatureConversion.getInputMapping(i)) {
+      // If inputMap->replacementValue is not nullptr, then the argument is
+      // dropped and a replacement value is provided to be the remappedValue.
+      if (inputMap->replacementValue) {
+        assert(inputMap->size == 0 &&
+               "invalid to provide a replacement value when the argument isn't "
+               "dropped");
+        remappedValues = inputMap->replacementValue;
+      } else
         remappedValues = newArgRef.slice(inputMap->inputNo, inputMap->size);
+    }
 
     BlockArgument *arg = block->getArgument(i);
     newArgMapping.push_back(convertArgument(arg, remappedValues, mapping));
@@ -1314,7 +1323,17 @@ void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo,
                                                     unsigned newInputCount) {
   assert(!remappedInputs[origInputNo] && "input has already been remapped");
   assert(newInputCount != 0 && "expected valid input count");
-  remappedInputs[origInputNo] = InputMapping{newInputNo, newInputCount};
+  remappedInputs[origInputNo] =
+      InputMapping{newInputNo, newInputCount, /*replacementValue=*/nullptr};
+}
+
+/// Remap an input of the original signature to another `replacementValue`
+/// value. This would make the signature converter drop this argument.
+void TypeConverter::SignatureConversion::remapInput(unsigned origInputNo,
+                                                    Value *replacementValue) {
+  assert(!remappedInputs[origInputNo] && "input has already been remapped");
+  remappedInputs[origInputNo] =
+      InputMapping{origInputNo, /*size=*/0, replacementValue};
 }
 
 /// This hooks allows for converting a type.