make raise activation use same getSplatConstantElementAttribute utils function

Signed-off-by: Matthew Francis-Landau <[email protected]>

Author: Matthew Francis-Landau

mlir-tensorrt/tensorrt/lib/TensorRT/Transforms/RaiseActivations.cpp

@@ -23,6 +23,7 @@
 //===----------------------------------------------------------------------===//
 #include "mlir-tensorrt-dialect/TensorRT/IR/TensorRTDialect.h"
 #include "mlir-tensorrt-dialect/TensorRT/Transforms/Passes.h"
+#include "mlir-tensorrt-dialect/TensorRT/Utils/Utils.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/IR/Matchers.h"

mlir-tensorrt/tensorrt/lib/TensorRT/Transforms/RaiseActivations.pdll

@@ -122,93 +122,22 @@ Constraint Erf(x: Value) -> Op {
 }
 
 Rewrite GetSplatElementAttr(x: Value) -> Attr [{
-  while(true) {
-    if(auto expandRank = x.getDefiningOp<tensorrt::ExpandRankOp>())
-      x = expandRank.getInput();
-    else if(auto reshape = x.getDefiningOp<tensorrt::ReshapeOp>())
-      x = reshape.getInput();
-    else if(auto broadcast = x.getDefiningOp<tensorrt::BroadcastOp>())
-      x = broadcast.getInput();
-    else if(auto cast = x.getDefiningOp<tensorrt::CastOp>())
-      x = cast.getInput();
-    else if(auto identity = x.getDefiningOp<tensorrt::IdentityOp>())
-      x = identity.getInput();
-    else if(auto slice = x.getDefiningOp<tensorrt::SliceOp>())
-      x = slice.getInput();
-    else if(auto constant = x.getDefiningOp<tensorrt::ConstantOp>()) {
-      DenseElementsAttr els{};
-      if(!matchPattern(x, m_Constant(&els)))
-        return {};
-      if(!els.isSplat())
-        return {};
-      Attribute value = els.getSplatValue<Attribute>();
-      return value;
-    } else
-      return {};
-  }
-  return {};
+  return *getSplatConstantElementAttribute(x);
 }];
 
 Constraint HasSplatElements(x: Value) [{
-  while(true) {
-    if(auto expandRank = x.getDefiningOp<tensorrt::ExpandRankOp>())
-      x = expandRank.getInput();
-    else if(auto reshape = x.getDefiningOp<tensorrt::ReshapeOp>())
-      x = reshape.getInput();
-    else if(auto broadcast = x.getDefiningOp<tensorrt::BroadcastOp>())
-      x = broadcast.getInput();
-    else if(auto cast = x.getDefiningOp<tensorrt::CastOp>())
-      x = cast.getInput();
-    else if(auto identity = x.getDefiningOp<tensorrt::IdentityOp>())
-      x = identity.getInput();
-    else if(auto slice = x.getDefiningOp<tensorrt::SliceOp>())
-      x = slice.getInput();
-    else if(auto constant = x.getDefiningOp<tensorrt::ConstantOp>()) {
-      DenseElementsAttr els{};
-      if(!matchPattern(x, m_Constant(&els)))
-        return failure();
-      if(!els.isSplat())
-        return failure();
-      Attribute value = els.getSplatValue<Attribute>();
-      return success(isa<FloatAttr>(value));
-    } else
-      return failure();
-  }
-  return failure();
+  return LogicalResult(getSplatConstantElementAttribute(x));
 }];
 
 /// Is true if `x` is a constant op that has a splat constant
 /// where splat element is equal to `attr`.
 Constraint SplatElements(x: Value, attr: Attr) [{
-  while(true) {
-    if(auto expandRank = x.getDefiningOp<tensorrt::ExpandRankOp>())
-      x = expandRank.getInput();
-    else if(auto reshape = x.getDefiningOp<tensorrt::ReshapeOp>())
-      x = reshape.getInput();
-    else if(auto broadcast = x.getDefiningOp<tensorrt::BroadcastOp>())
-      x = broadcast.getInput();
-    else if(auto cast = x.getDefiningOp<tensorrt::CastOp>())
-      x = cast.getInput();
-    else if(auto identity = x.getDefiningOp<tensorrt::IdentityOp>())
-      x = identity.getInput();
-    else if(auto slice = x.getDefiningOp<tensorrt::SliceOp>())
-      x = slice.getInput();
-    else if(auto constant = x.getDefiningOp<tensorrt::ConstantOp>()) {
-      DenseElementsAttr els{};
-      if(!matchPattern(x, m_Constant(&els)))
-        return failure();
-      if(!els.isSplat())
-        return failure();
-      Attribute value = els.getSplatValue<Attribute>();
-      if(!value) return failure();
-      if(value == attr) return success();
-      FloatAttr fvalue = dyn_cast<FloatAttr>(value);
-      FloatAttr fattr = dyn_cast<FloatAttr>(attr);
-      return success(fvalue && fattr && std::abs(fvalue.getValueAsDouble() - fattr.getValueAsDouble()) < .001); // handle different floating point type
-    } else
-      return failure();
-  }
-  return failure();
+  FailureOr<Attribute> value = getSplatConstantElementAttribute(x);
+  if(LogicalResult(value).failed()) return failure();
+  if(*value == attr) return success();
+  FloatAttr fvalue = dyn_cast<FloatAttr>(*value);
+  FloatAttr fattr = dyn_cast<FloatAttr>(attr);
+  return success(fvalue && fattr && std::abs(fvalue.getValueAsDouble() - fattr.getValueAsDouble()) < .001); // handle different floating point type
 }];
 
 /// We need a native C++ function since we can't create the right