[Mlir-commits] [mlir] a0b2104 - [mlir] Use TypedValue in single result traits

[Rahul Kayaith]

Author: Rahul Kayaith
Date: 2023-02-01T21:54:53-05:00
New Revision: a0b21046d8a3891b36fc992d97bc9daf4b3ce385

URL: https://github.com/llvm/llvm-project/commit/a0b21046d8a3891b36fc992d97bc9daf4b3ce385
DIFF: https://github.com/llvm/llvm-project/commit/a0b21046d8a3891b36fc992d97bc9daf4b3ce385.diff

LOG: [mlir] Use TypedValue in single result traits

Ops with a single result currently get a `getResult()` method +
conversion operator to `Value` through the `OneResult` trait. By moving
these to the `OneTypedResult` trait instead, we can use `TypedValue` as
the return type to get more specfic types.

When the result type is unknown ODS adds the
`OneTypedResult<mlir::Type>` trait, in which case there is no change in
the resulting API.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D142852

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpDefinition.h
    mlir/test/lib/Dialect/Test/TestPatterns.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 385597cc9a510..60c961f91cfac 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -601,17 +601,11 @@ struct MultiResultTraitBase : public TraitBase<ConcreteType, TraitType> {
 template <typename ConcreteType>
 class OneResult : public TraitBase<ConcreteType, OneResult> {
 public:
-  Value getResult() { return this->getOperation()->getResult(0); }
-
-  /// If the operation returns a single value, then the Op can be implicitly
-  /// converted to an Value. This yields the value of the only result.
-  operator Value() { return getResult(); }
-
   /// Replace all uses of 'this' value with the new value, updating anything
   /// in the IR that uses 'this' to use the other value instead.  When this
   /// returns there are zero uses of 'this'.
   void replaceAllUsesWith(Value newValue) {
-    getResult().replaceAllUsesWith(newValue);
+    this->getOperation()->getResult(0).replaceAllUsesWith(newValue);
   }
 
   /// Replace all uses of 'this' value with the result of 'op'.
@@ -637,10 +631,15 @@ class OneTypedResult {
   class Impl
       : public TraitBase<ConcreteType, OneTypedResult<ResultType>::Impl> {
   public:
-    ResultType getType() {
-      auto resultTy = this->getOperation()->getResult(0).getType();
-      return resultTy.template cast<ResultType>();
+    TypedValue<ResultType> getResult() {
+      return this->getOperation()->getResult(0);
     }
+
+    /// If the operation returns a single value, then the Op can be implicitly
+    /// converted to a Value. This yields the value of the only result.
+    operator TypedValue<ResultType>() { return getResult(); }
+
+    ResultType getType() { return getResult().getType(); }
   };
 };
 

diff  --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index de4fa39218844..a15d30725c1ad 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -579,7 +579,7 @@ struct TestUndoBlockArgReplace : public ConversionPattern {
     auto illegalOp =
         rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getF32Type());
     rewriter.replaceUsesOfBlockArgument(op->getRegion(0).getArgument(0),
-                                        illegalOp);
+                                        illegalOp->getResult(0));
     rewriter.updateRootInPlace(op, [] {});
     return success();
   }