[Mlir-commits] [mlir] [mlir][tosa] Optimize block scaled cast sequences (PR #188018)

[Sayan Saha]

================
@@ -935,6 +935,59 @@ void CastOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<NonNarrowingCastsOptimization>(context);
 }
 
+struct CancellingBlockScaledCastsOptimization
+    : public OpRewritePattern<tosa::CastToBlockScaledOp> {
+  using OpRewritePattern<tosa::CastToBlockScaledOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::CastToBlockScaledOp castToBlockScaledOp,
+                                PatternRewriter &rewriter) const override {
+    const Value castToBlockScaledInput = castToBlockScaledOp.getInputData();
+    auto castFromBlockScaledOp =
+        castToBlockScaledInput.getDefiningOp<tosa::CastFromBlockScaledOp>();
+    if (!castFromBlockScaledOp)
+      return rewriter.notifyMatchFailure(
+          castToBlockScaledOp,
+          "input must be cast_from_block_scaled operation");
+
+    const Value innerData = castFromBlockScaledOp.getInputData();
+    const Value innerScale = castFromBlockScaledOp.getInputScale();
+    const auto innerDataTy =
+        dyn_cast<ShapedType>(innerData.getType()).getElementType();
+    const auto innerScaleTy =
+        dyn_cast<ShapedType>(innerScale.getType()).getElementType();
+
+    const Value outerData = castToBlockScaledOp.getOutputData();
+    const Value outerScale = castToBlockScaledOp.getOutputScale();
+    const auto outerDataTy =
+        dyn_cast<ShapedType>(outerData.getType()).getElementType();
+    const auto outerScaleTy =
+        dyn_cast<ShapedType>(outerScale.getType()).getElementType();
----------------
sahas3 wrote:

these can be `cast` as per the ODS definition since the operands can only be `Tosa_TensorAtLeast1D`; alternatively there should be a `nullptr` check with `dyn_cast` for robustness

https://github.com/llvm/llvm-project/pull/188018