Merge pull request #411 from Xilinx/matthias.fix_tosa_pow

TorchToTosa: Correctly lower pow with broadcasting
Xilinx · Dec 12, 2024 · 2341f20 · 2341f20
2 parents c48183a + 0f604c8
commit 2341f20
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/lib/Conversion/TorchToTosa/TorchToTosa.cpp b/lib/Conversion/TorchToTosa/TorchToTosa.cpp
@@ -1084,24 +1084,26 @@ LogicalResult ConvertAtenOp<AtenPowTensorTensorOp>::matchAndRewrite(
     ConversionPatternRewriter &rewriter) const {
 
   Value self = adaptor.getSelf();
-  auto selfTy = cast<RankedTensorType>(self.getType());
+  auto selfTy = dyn_cast<RankedTensorType>(self.getType());
+  auto outType =
+      dyn_cast<RankedTensorType>(getTypeConverter()->convertType(op.getType()));
+  Value expTensor = adaptor.getExponent();
+  auto expTensorTy = dyn_cast<RankedTensorType>(expTensor.getType());
 
-  if (!selfTy)
+  if (!selfTy || !outType || !expTensorTy) {
     return rewriter.notifyMatchFailure(
         op, "Only ranked tensor types supported in TOSA Pow");
+  }
 
-  if (!isa<mlir::FloatType>(selfTy.getElementType()))
+  if (!isa<mlir::FloatType>(selfTy.getElementType())) {
     return rewriter.notifyMatchFailure(
         op, "Only floating-point datatype legalization supported");
+  }
 
-  auto outType =
-      cast<TensorType>(getTypeConverter()->convertType(op.getType()));
-
-  Value expTensor = adaptor.getExponent();
-  if (expTensor.getType() != selfTy) {
+  if (expTensorTy.getElementType() != selfTy.getElementType()) {
     expTensor = rewriter.createOrFold<tosa::CastOp>(
         op->getLoc(),
-        RankedTensorType::get(outType.getShape(), selfTy.getElementType()),
+        RankedTensorType::get(expTensorTy.getShape(), selfTy.getElementType()),
         expTensor);
   }
 

diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -1993,6 +1993,7 @@
     "ElementwisePowTensorBroadcastModule_basic",
     "ElementwisePowTensorBroadcastStaticModule_basic",
     "ElementwisePowTensorModule_basic",
+    "ElementwisePowTensorStaticModule_basic",
     "ElementwisePreluModule_basic",
     "ElementwisePreluStaticModule_basic",
     "ElementwiseRad2DegModule_basic",