Merge pull request #454 from Xilinx/jrickert.mixed_dialects

Adjust IndexExprBuilderForAnalysis::getConst to not only handle onnx.…

Author: jorickert

src/Dialect/ONNX/DialectBuilder.cpp

@@ -887,7 +887,7 @@ Value OnnxBuilder::getOrCastToI8(Value val, bool simpleCast) {
 
 // Return null if none is found.
 ElementsAttr IndexExprBuilderForAnalysis::getConst(Value value) {
-  return getElementAttributeFromONNXValue(value);
+  return getElementAttributeFromConstLikeValue(value);
 }
 
 // Return null if the value at index i is not a constant.

src/Dialect/ONNX/ONNXOps/OpHelper.cpp

@@ -316,6 +316,29 @@ ElementsAttr getElementAttributeFromONNXValue(Value value) {
   return nullptr;
 }
 
+ElementsAttr getElementAttributeFromConstLikeValue(Value value) {
+  auto *definingOp = value.getDefiningOp();
+  if (!isConstLikeOperation(definingOp)) {
+    return nullptr;
+  }
+  SmallVector<OpFoldResult, 1> foldResults;
+  [[maybe_unused]] const LogicalResult folded = definingOp->fold(foldResults);
+  assert(succeeded(folded) && "ConstantLike op failed to fold");
+  assert(foldResults.size() == 1 &&
+         "ConstantLike op fold produced more results than expected");
+  auto foldAttr = dyn_cast<Attribute>(foldResults[0]);
+  assert(foldAttr && "ConstantLike op fold did not return an Attribute");
+  return dyn_cast<ElementsAttr>(foldAttr);
+}
+
+bool isConstLikeValue(Value value) {
+  return isConstLikeOperation(value.getDefiningOp());
+}
+
+bool isConstLikeOperation(mlir::Operation *op) {
+  return op && op->hasTrait<OpTrait::ConstantLike>();
+}
+
 // compare two ElementsAttr, except for their internal buffer size
 bool compareValueFromElementAttribute(
     ElementsAttr &attr1, ElementsAttr &attr2) {

src/Dialect/ONNX/ONNXOps/OpHelper.hpp

@@ -174,6 +174,14 @@ void ArrayAttrIntVals(mlir::ArrayAttr a, mlir::SmallVectorImpl<int64_t> &i);
 
 mlir::ElementsAttr getElementAttributeFromONNXValue(mlir::Value value);
 
+// Get a ElementsAttr from a value that is defined by a ConstantLike op that
+// folds to an ElementsAttr
+mlir::ElementsAttr getElementAttributeFromConstLikeValue(mlir::Value value);
+
+[[nodiscard]] bool isConstLikeValue(mlir::Value value);
+
+[[nodiscard]] bool isConstLikeOperation(mlir::Operation *op);
+
 bool compareValueFromElementAttribute(
     mlir::ElementsAttr &attr1, mlir::ElementsAttr &attr2);
 

src/Dialect/ONNX/ONNXOps/Tensor/Slice.cpp

@@ -160,7 +160,7 @@ LogicalResult ONNXSliceOp::inferShapes(
 
   // Cannot infer shape if axes is not a constant. It can be a constant after
   // several rounds of shape-inference and constant propagation.
-  if (!isNoneValue(axes) && !getONNXConstantOp(axes))
+  if (!isNoneValue(axes) && !isConstLikeValue(axes))
     return success();
 
   const auto startsType =

test/mlir/onnx/onnx_shape_inference.mlir

@@ -4454,4 +4454,17 @@ func.func @test_slice_negative_steps(%arg0: tensor<100x200xf32>) -> tensor<*xf32
 }
 // CHECK-LABEL:  func.func @test_slice_negative_steps
 // CHECK:          "onnx.Slice"
-// CHECK-SAME:       (tensor<100x200xf32>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<16x16xf32>
+// CHECK-SAME:       (tensor<100x200xf32>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<16x16xf32>
+
+// -----
+func.func @test_slice_negative_steps_mixed_dialects(%arg0: tensor<100x200xf32>) -> tensor<*xf32> {
+  %axes = "tosa.const"() {value = dense<[0, 1]> : tensor<2xi64> } : () -> tensor<2xi64>
+  %starts = "tosa.const"() {value = dense<[-10, -20]> : tensor<2xi64> } : () -> tensor<2xi64>
+  %ends = "tosa.const"() {value = dense<[10, 20]> : tensor<2xi64> } : () -> tensor<2xi64>
+  %steps = "tosa.const"() {value = dense<[-5, -10]> : tensor<2xi64> } : () -> tensor<2xi64>
+  %1 = "onnx.Slice"(%arg0, %starts, %ends, %axes, %steps) : (tensor<100x200xf32>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<*xf32>
+  return %1 : tensor<*xf32> 
+}
+// CHECK-LABEL:  func.func @test_slice_negative_steps_mixed_dialects
+// CHECK:          "onnx.Slice"
+// CHECK-SAME:       (tensor<100x200xf32>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<16x16xf32>