Repository: iree-org/iree
Pull request #20607: [GlobalOptimizations] Add a pass to simplify strided contraction-like ops
Author: login: zjgarvey
Created: 2025-04-22T21:52:42Z | Merged: 2025-04-24T18:34:01Z

This PR adds a pattern to manipulate generic ops which satisfy all "contraction" conditions except for the indexing maps being projected permutations. Namely, if the input indexing map has results of the form `dim * cst`, this pattern will factor the original generic op into `tensor.extract_slice + contraction linalg.generic`.

This addresses #20600 . Each of the included lit test examples were not compiling to mfma instructions before this patch.


Diff:
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
index b0c24993226b..cc5c8612cfa5 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
+++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -46,6 +46,7 @@ iree_compiler_cc_library(
     srcs = [
         "CleanupNumericNarrowing.cpp",
         "Convert1X1FilterConv2DToMatmul.cpp",
+        "ConvertStridedContractionToContraction.cpp",
         "DataLayoutPropagation.cpp",
         "DecomposeConcat.cpp",
         "DemoteContractionInputsToBF16.cpp",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
index b07be2cb3f3e..50fe8f77a01f 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
+++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -42,6 +42,7 @@ iree_cc_library(
   SRCS
     "CleanupNumericNarrowing.cpp"
     "Convert1X1FilterConv2DToMatmul.cpp"
+    "ConvertStridedContractionToContraction.cpp"
     "DataLayoutPropagation.cpp"
     "DecomposeConcat.cpp"
     "DemoteContractionInputsToBF16.cpp"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp b/compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
new file mode 100644
index 000000000000..bf4de1081ea9
--- /dev/null
+++ b/compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
@@ -0,0 +1,145 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            })) {
+      return failure();
+    }
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    unsigned inputPos = op.getDpsInputOperand(0)->getOperandNumber();
+    unsigned filterPos = op.getDpsInputOperand(1)->getOperandNumber();
+    unsigned resInitPos = op.getDpsInitOperand(0)->getOperandNumber();
+    AffineMap inputMap = mapRange[inputPos];
+    AffineMap filterMap = mapRange[filterPos];
+    AffineMap resultMap = mapRange[resInitPos];
+    // For now, we are only handling the case where the first input is the
+    // only non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation()) {
+      return failure();
+    }
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    SmallVector<int64_t, 4> staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    // Walk through input map and look for expressions of the form `dim * cst`.
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // Skip dim exprs and constant exprs.
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // Look at binary op expressions.
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // Fail if we see some unexpected kind of expression.
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // Binary expressions must be of the form `dim * cst`.
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
+      inputShape[pos] = newSize;
+      replacementExprs.push_back(lhs);
+    }
+
+    // Fail if we don't have any work to do.
+    if (strides.empty())
+      return failure();
+
+    mapRange[inputPos] =
+        AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(),
+                       replacementExprs, op.getContext());
+    auto sliceTy = RankedTensorType::get(inputShape, inputTy.getElementType());
+
+    unsigned rank = inputTy.getRank();
+    SmallVector<OpFoldResult> vOffset(rank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> vSizes;
+    SmallVector<OpFoldResult> vStride(rank, rewriter.getIndexAttr(1));
+    Location loc = op.getLoc();
+    for (unsigned i = 0; i < inputTy.getRank(); i++) {
+      if (strides.contains(i)) {
+        vStride[i] = rewriter.getIndexAttr(strides.at(i));
+      }
+      if (inputShape[i] != ShapedType::kDynamic) {
+        vSizes.push_back(rewriter.getIndexAttr(inputShape[i]));
+        continue;
+      }
+      vSizes.push_back(rewriter.createOrFold<tensor::DimOp>(loc, input, i));
+    }
+    Value extractedSlice = rewriter.create<tensor::ExtractSliceOp>(
+        loc, sliceTy, input, vOffset, vSizes, vStride);
+    rewriter.startOpModification(op);
+    op.setIndexingMapsAttr(rewriter.getAffineMapArrayAttr(mapRange));
+    op.setOperand(0, extractedSlice);
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+struct ConvertStridedContractionToContractionPass
+    : public impl::ConvertStridedContractionToContractionPassBase<
+          ConvertStridedContractionToContractionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithDialect, tensor::TensorDialect>();
+  }
+
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(&getContext());
+    patterns.insert<ConvertStridedContractionToContraction>(context);
+    walkAndApplyPatterns(getOperation(), std::move(patterns));
+  }
+};
+} // namespace
+} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
index 8a9843acae5b..c4dae945c4c5 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
+++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -155,6 +155,8 @@ void buildGlobalOptimizationPassPipeline(
       });
 
   mainPassManager.addPass(DispatchCreation::createFoldUnitExtentDimsPass());
+  mainPassManager.addPass(
+      GlobalOptimization::createConvertStridedContractionToContractionPass());
   FunctionLikeNest(mainPassManager)
       .addPredicatedPass(clEnableFuseSiluHorizontalMatmul,
                          createFuseSiluHorizontalMatmulPass)
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
index 694ff5faf467..d7fe5554284f 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/Passes.td
+++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -19,6 +19,11 @@ def Convert1X1FilterConv2DToMatmulPass:
   let summary = "Convert linalg convolution ops with 1x1 kernels into linalg matrix multiplication ops.";
 }
 
+def ConvertStridedContractionToContractionPass:
+    Pass<"iree-global-opt-convert-strided-contraction-to-contraction", ""> {
+  let summary = "Factors out an extract_slice from contraction-like ops with strided inputs.";
+}
+
 def DecomposeConcatPass :
     Pass<"iree-global-opt-decompose-concat", ""> {
   let summary = "Decomposes concatenations into a destination and a sequence of slice inserts.";
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
index 2b6d926ab3ee..dcdacfc12188 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -34,6 +34,7 @@ iree_lit_test_suite(
             "propagate_linalg_transpose.mlir",
             "raise_special_ops.mlir",
             "remove_zero_extent_tensors.mlir",
+            "strided_contraction_to_contraction.mlir",
             "transformation_pipeline.mlir",
             "transpose_and_decompose_concat.mlir",
             "warn_on_uninitialized_values.mlir",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
index b291505aaba7..c8f50e8add40 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -32,6 +32,7 @@ iree_lit_test_suite(
     "propagate_linalg_transpose.mlir"
     "raise_special_ops.mlir"
     "remove_zero_extent_tensors.mlir"
+    "strided_contraction_to_contraction.mlir"
     "transformation_pipeline.mlir"
     "transpose_and_decompose_concat.mlir"
     "warn_on_uninitialized_values.mlir"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/strided_contraction_to_contraction.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/strided_contraction_to_contraction.mlir
new file mode 100644
index 000000000000..780bfbf4d22f
--- /dev/null
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/strided_contraction_to_contraction.mlir
@@ -0,0 +1,117 @@
+// RUN: iree-opt --split-input-file --mlir-print-local-scope -iree-global-opt-convert-strided-contraction-to-contraction %s | FileCheck %s
+
+util.func public @strided_from_output_static(%input: tensor<2x118x182x448xbf16>, %filter: tensor<896x448xbf16>) -> tensor<2x59x91x896xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<2x59x91x896xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<2x59x91x896xf32>) -> tensor<2x59x91x896xf32>
+  %2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d0, 2 * d1, d2 * 2, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%input, %filter : tensor<2x118x182x448xbf16>, tensor<896x448xbf16>) outs(%1 : tensor<2x59x91x896xf32>) {
+  ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+    %3 = arith.extf %in : bf16 to f32
+    %4 = arith.extf %in_0 : bf16 to f32
+    %5 = arith.mulf %3, %4 : f32
+    %6 = arith.addf %out, %5 : f32
+    linalg.yield %6 : f32
+  } -> tensor<2x59x91x896xf32>
+  util.return %2 : tensor<2x59x91x896xf32>
+}
+
+// CHECK-LABEL: @strided_from_output_static(
+// CHECK-SAME:      %[[INPUT:.*]]: tensor<2x118x182x448xbf16>
+// CHECK-SAME:      %[[FILTER:.*]]: tensor<896x448xbf16>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[INPUT]][0, 0, 0, 0] [2, 59, 91, 448] [1, 2, 2, 1]
+// CHECK-SAME:      tensor<2x118x182x448xbf16> to tensor<2x59x91x448xbf16>
+// CHECK: %[[GEN:.*]] = linalg.generic
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d4)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>
+// CHECK-SAME:      ins(%[[SLICE]], %[[FILTER]]
+// CHECK: util.return %[[GEN]]
+
+
+// -----
+
+util.func public @strided_from_output_dynamic_batch(%input: tensor<?x118x182x448xbf16>, %filter: tensor<896x448xbf16>) -> tensor<?x59x91x896xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %dim = tensor.dim %input, %c0 : tensor<?x118x182x448xbf16>
+  %0 = tensor.empty(%dim) : tensor<?x59x91x896xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x59x91x896xf32>) -> tensor<?x59x91x896xf32>
+  %2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d0, d1 * 2, d2 * 2, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%input, %filter : tensor<?x118x182x448xbf16>, tensor<896x448xbf16>) outs(%1 : tensor<?x59x91x896xf32>) {
+  ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+    %3 = arith.extf %in : bf16 to f32
+    %4 = arith.extf %in_0 : bf16 to f32
+    %5 = arith.mulf %3, %4 : f32
+    %6 = arith.addf %out, %5 : f32
+    linalg.yield %6 : f32
+  } -> tensor<?x59x91x896xf32>
+  util.return %2 : tensor<?x59x91x896xf32>
+}
+
+// CHECK-LABEL: @strided_from_output_dynamic_batch(
+// CHECK-SAME:      %[[INPUT:.*]]: tensor<?x118x182x448xbf16>
+// CHECK-SAME:      %[[FILTER:.*]]: tensor<896x448xbf16>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[INPUT]][0, 0, 0, 0] [%[[DIM:.*]], 59, 91, 448] [1, 2, 2, 1]
+// CHECK-SAME:      tensor<?x118x182x448xbf16> to tensor<?x59x91x448xbf16>
+// CHECK: %[[GEN:.*]] = linalg.generic
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d4)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>
+// CHECK-SAME:      ins(%[[SLICE]], %[[FILTER]]
+// CHECK: util.return %[[GEN]]
+
+// -----
+
+util.func public @strided_from_output_partial_conv(%input: tensor<2x118x182x448xbf16>, %filter: tensor<896x2x448xbf16>) -> tensor<2x59x91x896xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<2x59x91x896xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<2x59x91x896xf32>) -> tensor<2x59x91x896xf32>
+  %2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2, d2 * 2 + d4, d5)>, affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>, affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<2x118x182x448xbf16>, tensor<896x2x448xbf16>) outs(%1 : tensor<2x59x91x896xf32>) {
+  ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+    %3 = arith.extf %in : bf16 to f32
+    %4 = arith.extf %in_0 : bf16 to f32
+    %5 = arith.mulf %3, %4 : f32
+    %6 = arith.addf %out, %5 : f32
+    linalg.yield %6 : f32
+  } -> tensor<2x59x91x896xf32>
+  util.return %2 : tensor<2x59x91x896xf32>
+}
+
+// CHECK-LABEL: @strided_from_output_partial_conv
+// CHECK-SAME:      %[[INPUT:.*]]: tensor<2x118x182x448xbf16>
+// CHECK-SAME:      %[[FILTER:.*]]: tensor<896x2x448xbf16>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[INPUT]][0, 0, 0, 0] [2, 59, 182, 448] [1, 2, 1, 1]
+// CHECK-SAME:     tensor<2x118x182x448xbf16> to tensor<2x59x182x448xbf16>
+// CHECK: linalg.generic
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2 * 2 + d4, d5)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+// CHECK-SAME:      ins(%[[SLICE]], %[[FILTER]]
+
+// -----
+
+util.func public @strided_from_filter_static(%input: tensor<896x118x16xbf16>, %filter: tensor<448x59x16xbf16>) -> tensor<896x448xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<896x448xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<896x448xf32>) -> tensor<896x448xf32>
+  %2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d2 * 2, d3)>, affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%input, %filter : tensor<896x118x16xbf16>, tensor<448x59x16xbf16>) outs(%1 : tensor<896x448xf32>) {
+  ^bb0(%in: bf16, %in_0: bf16, %out: f32):
+    %3 = arith.extf %in : bf16 to f32
+    %4 = arith.extf %in_0 : bf16 to f32
+    %5 = arith.mulf %3, %4 : f32
+    %6 = arith.addf %out, %5 : f32
+    linalg.yield %6 : f32
+  } -> tensor<896x448xf32>
+  util.return %2 : tensor<896x448xf32>
+}
+
+// CHECK-LABEL: @strided_from_filter_static(
+// CHECK-SAME:      %[[INPUT:.*]]: tensor<896x118x16xbf16>
+// CHECK-SAME:      %[[FILTER:.*]]: tensor<448x59x16xbf16>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[INPUT]][0, 0, 0] [896, 59, 16] [1, 2, 1]
+// CHECK-SAME:      tensor<896x118x16xbf16> to tensor<896x59x16xbf16>
+// CHECK: %[[GEN:.*]] = linalg.generic
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>
+// CHECK-SAME:      affine_map<(d0, d1, d2, d3) -> (d0, d1)>
+// CHECK-SAME:      ins(%[[SLICE]], %[[FILTER]]
+// CHECK: util.return %[[GEN]]
diff --git a/compiler/src/iree/compiler/Preprocessing/Passes.cpp b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
index 89b754112f39..ec888cbffe42 100644
--- a/compiler/src/iree/compiler/Preprocessing/Passes.cpp
+++ b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
@@ -147,6 +147,8 @@ buildMakeSingleDispatchPassPipeline(OpPassManager &passManager,
   // Generalize transposes and any other remaining named linalg ops that can
   // now be represented as generics.
   passManager.addPass(GlobalOptimization::createGeneralizeLinalgNamedOpsPass());
+  passManager.addPass(
+      GlobalOptimization::createConvertStridedContractionToContractionPass());
   passManager.addPass(DispatchCreation::createFusionPreprocessingPass());
   passManager.addPass(mlir::createCSEPass());
   DispatchCreation::BubbleUpExpandShapesPassOptions bubbleOptions;


Reviews:
user: login: nirvedhmeshram
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T22:04:58Z
user: login: zjgarvey
body: Some self-review and questions.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T16:35:11Z
user: login: zjgarvey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T16:50:24Z
user: login: MaheshRavishankar
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T17:38:48Z
user: login: zjgarvey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T20:28:00Z
user: login: zjgarvey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T20:31:43Z
user: login: qedawkins
body: Some nits. Pass location within GlobalOptimization makes sense though and otherwise LG
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T21:31:30Z
user: login: zjgarvey
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T22:14:58Z
user: login: qedawkins
body: LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T16:31:15Z

Review comments:
diff_hunk: @@ -147,6 +147,8 @@ buildMakeSingleDispatchPassPipeline(OpPassManager &passManager,
   // Generalize transposes and any other remaining named linalg ops that can
   // now be represented as generics.
   passManager.addPass(GlobalOptimization::createGeneralizeLinalgNamedOpsPass());
+  passManager.addPass(
path: compiler/src/iree/compiler/Preprocessing/Passes.cpp
user: login: nirvedhmeshram
body: I would also put the pass in src/iree/compiler/GlobalOptimization/Passes.cpp after the first `createGeneralizeLinalgNamedOpsPass` pass as well as that way the CI can catch any missed issues and have the general compiler benefit from it too.
created_at: 2025-04-22T22:04:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: Are we always guaranteed to have the `outs` at the back of this array?
created_at: 2025-04-23T16:24:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form cst*dim
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: I've added this in case to try and support when only one of the complex dims is a `dim * cst` expression, but the other is convolution-like. The last lit test is an example of this (and was previously not going down the IGEMM convolution pipeline). 
created_at: 2025-04-23T16:27:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form cst*dim
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      noReplacement = false;
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: Would it be beneficial to support dynamic strided dims?
created_at: 2025-04-23T16:28:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form cst*dim
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      noReplacement = false;
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
+      inputShape[pos] = newSize;
+      replacementExprs.push_back(lhs);
+    }
+
+    // fail if we don't have any work to do
+    if (noReplacement)
+      return failure();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: I can probably remove the bool and check if `strides.empty()`
created_at: 2025-04-23T16:31:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: When there is a `mul` between a `AffineConstantExpr` and `AffineDimExpr`, is the constant always expected to be on the right?
created_at: 2025-04-23T16:32:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -147,6 +147,8 @@ buildMakeSingleDispatchPassPipeline(OpPassManager &passManager,
   // Generalize transposes and any other remaining named linalg ops that can
   // now be represented as generics.
   passManager.addPass(GlobalOptimization::createGeneralizeLinalgNamedOpsPass());
+  passManager.addPass(
path: compiler/src/iree/compiler/Preprocessing/Passes.cpp
user: login: zjgarvey
body: I added it after `fold-unit-extent-dims`, since the unit dim folding is likely where we would get conv-like generic ops reducing to strided contractions.
created_at: 2025-04-23T16:50:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: MaheshRavishankar
body: THe way to get that would be `getMatchingIndexingMap`. So you can just get the `lhs`, `rhs` and `outs` `OpOperand` 
 and get indexing map from there
created_at: 2025-04-23T17:35:02Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: MaheshRavishankar
body: No that isnt always gauranteed. You can check for either lhs or rhs being the constant. 
created_at: 2025-04-23T17:36:15Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form cst*dim
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      noReplacement = false;
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
+      inputShape[pos] = newSize;
+      replacementExprs.push_back(lhs);
+    }
+
+    // fail if we don't have any work to do
+    if (noReplacement)
+      return failure();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: MaheshRavishankar
body: Yes, but either works fine.
created_at: 2025-04-23T17:37:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: I opted to get the operand numbers for the dpsinputs/dpsinits and use those since I need to modify the `mapRange` at the input position later anyway.
created_at: 2025-04-23T20:28:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: I changed a lit test so that one of the examples had `2 * dim` in the indexing map, and it still seems to swap the order to `dim * cst`, and the pattern applies successfully. I'm not sure if this is just a parsing feature. 

I have some code to check both ways, but the code is a bit uglier so I wasn't sure whether to include it. 
created_at: 2025-04-23T20:31:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: style nit: Add braces after multi-line if
created_at: 2025-04-23T20:48:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: Not a parsing feature, AffineMap::simplify will do this for you for both + and * operations with constants. Map simplification gets called pretty frequently so it probably happened somewhere in here.
created_at: 2025-04-23T20:50:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: Does `mlir::linalg::inferConvolutionDims` not work here? It returns a struct enumerating dim types and tries to capture strides + dilations as well, but might not recognize strides without the "convolution" expression.
created_at: 2025-04-23T20:53:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    unsigned inputPos = op.getDpsInputOperand(0)->getOperandNumber();
+    unsigned filterPos = op.getDpsInputOperand(1)->getOperandNumber();
+    unsigned resInitPos = op.getDpsInitOperand(0)->getOperandNumber();
+    AffineMap inputMap = mapRange[inputPos];
+    AffineMap filterMap = mapRange[filterPos];
+    AffineMap resultMap = mapRange[resInitPos];
+    // For now, we are only handling the case where the first input is the
+    // only non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: ```suggestion
    SmallVector<int64_t, 4> staticShape = op.getStaticLoopRanges();
```
created_at: 2025-04-23T20:53:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    unsigned inputPos = op.getDpsInputOperand(0)->getOperandNumber();
+    unsigned filterPos = op.getDpsInputOperand(1)->getOperandNumber();
+    unsigned resInitPos = op.getDpsInitOperand(0)->getOperandNumber();
+    AffineMap inputMap = mapRange[inputPos];
+    AffineMap filterMap = mapRange[filterPos];
+    AffineMap resultMap = mapRange[resInitPos];
+    // For now, we are only handling the case where the first input is the
+    // only non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    // walk through input map and look for expressions of the form `dim * cst`
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: Capitalize and add punctuation here and below.
```suggestion
    // Walk through input map and look for expressions of the form `dim * cst`.
```
created_at: 2025-04-23T20:54:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form cst*dim
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      noReplacement = false;
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: Support for something is always better than not, but you don't have to implement everything in one go
created_at: 2025-04-23T20:58:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    unsigned inputPos = op.getDpsInputOperand(0)->getOperandNumber();
+    unsigned filterPos = op.getDpsInputOperand(1)->getOperandNumber();
+    unsigned resInitPos = op.getDpsInitOperand(0)->getOperandNumber();
+    AffineMap inputMap = mapRange[inputPos];
+    AffineMap filterMap = mapRange[filterPos];
+    AffineMap resultMap = mapRange[resInitPos];
+    // For now, we are only handling the case where the first input is the
+    // only non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    // walk through input map and look for expressions of the form `dim * cst`
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form `dim * cst`
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
+      inputShape[pos] = newSize;
+      replacementExprs.push_back(lhs);
+    }
+
+    // fail if we don't have any work to do
+    if (strides.empty())
+      return failure();
+
+    mapRange[inputPos] =
+        AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(),
+                       replacementExprs, op.getContext());
+    auto sliceTy = RankedTensorType::get(inputShape, inputTy.getElementType());
+
+    unsigned rank = inputTy.getRank();
+    SmallVector<OpFoldResult> vOffset(rank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> vSizes;
+    SmallVector<OpFoldResult> vStride(rank, rewriter.getIndexAttr(1));
+    Location loc = op.getLoc();
+    for (unsigned i = 0; i < inputTy.getRank(); i++) {
+      if (strides.contains(i)) {
+        vStride[i] = rewriter.getIndexAttr(strides.at(i));
+      }
+      if (inputShape[i] != ShapedType::kDynamic) {
+        vSizes.push_back(rewriter.getIndexAttr(inputShape[i]));
+        continue;
+      }
+      vSizes.push_back(rewriter.createOrFold<tensor::DimOp>(loc, input, i));
+    }
+    Value extractedSlice = rewriter.create<tensor::ExtractSliceOp>(
+        loc, sliceTy, input, vOffset, vSizes, vStride);
+    rewriter.startOpModification(op);
+    op.setIndexingMapsAttr(rewriter.getAffineMapArrayAttr(mapRange));
+    op.setOperand(0, extractedSlice);
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+struct ConvertStridedContractionToContractionPass
+    : public impl::ConvertStridedContractionToContractionPassBase<
+          ConvertStridedContractionToContractionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect>();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: Since this pass operates on linalg ops, it doesn't need to depend on it.
created_at: 2025-04-23T21:20:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    unsigned inputPos = op.getDpsInputOperand(0)->getOperandNumber();
+    unsigned filterPos = op.getDpsInputOperand(1)->getOperandNumber();
+    unsigned resInitPos = op.getDpsInitOperand(0)->getOperandNumber();
+    AffineMap inputMap = mapRange[inputPos];
+    AffineMap filterMap = mapRange[filterPos];
+    AffineMap resultMap = mapRange[resInitPos];
+    // For now, we are only handling the case where the first input is the
+    // only non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    // walk through input map and look for expressions of the form `dim * cst`
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form `dim * cst`
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
+      inputShape[pos] = newSize;
+      replacementExprs.push_back(lhs);
+    }
+
+    // fail if we don't have any work to do
+    if (strides.empty())
+      return failure();
+
+    mapRange[inputPos] =
+        AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(),
+                       replacementExprs, op.getContext());
+    auto sliceTy = RankedTensorType::get(inputShape, inputTy.getElementType());
+
+    unsigned rank = inputTy.getRank();
+    SmallVector<OpFoldResult> vOffset(rank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> vSizes;
+    SmallVector<OpFoldResult> vStride(rank, rewriter.getIndexAttr(1));
+    Location loc = op.getLoc();
+    for (unsigned i = 0; i < inputTy.getRank(); i++) {
+      if (strides.contains(i)) {
+        vStride[i] = rewriter.getIndexAttr(strides.at(i));
+      }
+      if (inputShape[i] != ShapedType::kDynamic) {
+        vSizes.push_back(rewriter.getIndexAttr(inputShape[i]));
+        continue;
+      }
+      vSizes.push_back(rewriter.createOrFold<tensor::DimOp>(loc, input, i));
+    }
+    Value extractedSlice = rewriter.create<tensor::ExtractSliceOp>(
+        loc, sliceTy, input, vOffset, vSizes, vStride);
+    rewriter.startOpModification(op);
+    op.setIndexingMapsAttr(rewriter.getAffineMapArrayAttr(mapRange));
+    op.setOperand(0, extractedSlice);
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+struct ConvertStridedContractionToContractionPass
+    : public impl::ConvertStridedContractionToContractionPassBase<
+          ConvertStridedContractionToContractionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect>();
+    registry.insert<tensor::TensorDialect>();
+    registry.insert<arith::ArithDialect>();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: ```suggestion
    registry.insert<arith::ArithDialect, tensor::TensorDialect>();
```
created_at: 2025-04-23T21:20:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    unsigned inputPos = op.getDpsInputOperand(0)->getOperandNumber();
+    unsigned filterPos = op.getDpsInputOperand(1)->getOperandNumber();
+    unsigned resInitPos = op.getDpsInitOperand(0)->getOperandNumber();
+    AffineMap inputMap = mapRange[inputPos];
+    AffineMap filterMap = mapRange[filterPos];
+    AffineMap resultMap = mapRange[resInitPos];
+    // For now, we are only handling the case where the first input is the
+    // only non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    // walk through input map and look for expressions of the form `dim * cst`
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form `dim * cst`
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
+      inputShape[pos] = newSize;
+      replacementExprs.push_back(lhs);
+    }
+
+    // fail if we don't have any work to do
+    if (strides.empty())
+      return failure();
+
+    mapRange[inputPos] =
+        AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(),
+                       replacementExprs, op.getContext());
+    auto sliceTy = RankedTensorType::get(inputShape, inputTy.getElementType());
+
+    unsigned rank = inputTy.getRank();
+    SmallVector<OpFoldResult> vOffset(rank, rewriter.getIndexAttr(0));
+    SmallVector<OpFoldResult> vSizes;
+    SmallVector<OpFoldResult> vStride(rank, rewriter.getIndexAttr(1));
+    Location loc = op.getLoc();
+    for (unsigned i = 0; i < inputTy.getRank(); i++) {
+      if (strides.contains(i)) {
+        vStride[i] = rewriter.getIndexAttr(strides.at(i));
+      }
+      if (inputShape[i] != ShapedType::kDynamic) {
+        vSizes.push_back(rewriter.getIndexAttr(inputShape[i]));
+        continue;
+      }
+      vSizes.push_back(rewriter.createOrFold<tensor::DimOp>(loc, input, i));
+    }
+    Value extractedSlice = rewriter.create<tensor::ExtractSliceOp>(
+        loc, sliceTy, input, vOffset, vSizes, vStride);
+    rewriter.startOpModification(op);
+    op.setIndexingMapsAttr(rewriter.getAffineMapArrayAttr(mapRange));
+    op.setOperand(0, extractedSlice);
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+struct ConvertStridedContractionToContractionPass
+    : public impl::ConvertStridedContractionToContractionPassBase<
+          ConvertStridedContractionToContractionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect>();
+    registry.insert<tensor::TensorDialect>();
+    registry.insert<arith::ArithDialect>();
+  }
+
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(&getContext());
+    patterns.insert<ConvertStridedContractionToContraction>(context);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: qedawkins
body: We don't need greedy application here. Try `walkAndApplyPatterns` instead
created_at: 2025-04-23T21:21:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -155,6 +155,8 @@ void buildGlobalOptimizationPassPipeline(
       });
 
   mainPassManager.addPass(DispatchCreation::createFoldUnitExtentDimsPass());
+  mainPassManager.addPass(
+      GlobalOptimization::createConvertStridedContractionToContractionPass());
path: compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
user: login: qedawkins
body: This pass location is fine. It should also work in the backend but either way is ok.
created_at: 2025-04-23T21:23:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,147 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: I looked into that originally, but it ended up not working well when the filter dims are already folded out. See, in particular, <https://github.com/llvm/llvm-project/blob/ff6a23d65eccba0991244c4127506452e1030393/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp#L626>. I'm sure it might be better to try and put this somewhere before that folding happens, but I figured it would be better to write something that would work even if the generic didn't original come from a unit-filter-dim conv. 
created_at: 2025-04-23T21:55:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,146 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTSTRIDEDCONTRACTIONTOCONTRACTIONPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+class ConvertStridedContractionToContraction
+    : public OpRewritePattern<linalg::GenericOp> {
+public:
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(linalg::GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Check if the generic op satisfies all other conditions for being a
+    // contraction.
+    if (op.getNumDpsInputs() != 2 || op.getNumDpsInits() != 1)
+      return failure();
+    if (op.getNumReductionLoops() == 0)
+      return failure();
+    if (!mlir::linalg::detail::isContractionBody(
+            *op.getBlock(), [](Operation *first, Operation *second) {
+              if ((isa<arith::MulFOp>(first) && isa<arith::AddFOp>(second)) ||
+                  (isa<arith::MulIOp>(first) && isa<arith::AddIOp>(second)))
+                return true;
+              return false;
+            }))
+      return failure();
+
+    SmallVector<AffineMap> mapRange = op.getIndexingMapsArray();
+    AffineMap inputMap = mapRange[0];
+    AffineMap filterMap = mapRange[1];
+    AffineMap resultMap = mapRange[2];
+    // For now, we are only handling the case where the input is the only
+    // non-projected permutation.
+    if (!filterMap.isProjectedPermutation() ||
+        !resultMap.isProjectedPermutation())
+      return failure();
+    if (inputMap.isProjectedPermutation())
+      return failure();
+    auto staticShape = op.getStaticLoopRanges();
+
+    llvm::SmallDenseMap<unsigned, int64_t> strides;
+    SmallVector<AffineExpr> replacementExprs;
+    Value input = op.getDpsInputs()[0];
+    auto inputTy = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputTy)
+      return failure();
+    SmallVector<int64_t> inputShape(inputTy.getShape());
+    replacementExprs.reserve(inputMap.getNumResults());
+    bool noReplacement = true;
+    // walk through input map and look for expressions of the form cst*dim
+    for (auto [pos, expr] : llvm::enumerate(inputMap.getResults())) {
+      // skip dim exprs and constant exprs
+      if (isa<AffineDimExpr>(expr) || isa<AffineConstantExpr>(expr)) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      // look at binary op expressions
+      auto binexpr = dyn_cast<AffineBinaryOpExpr>(expr);
+      // fail if we see some unexpected kind of expression
+      if (!binexpr)
+        return failure();
+      auto rhs = dyn_cast<AffineConstantExpr>(binexpr.getRHS());
+      auto lhs = dyn_cast<AffineDimExpr>(binexpr.getLHS());
+      // binary expressions must be of the form cst*dim
+      if (!rhs || !lhs || binexpr.getKind() != AffineExprKind::Mul) {
+        replacementExprs.push_back(expr);
+        continue;
+      }
+      noReplacement = false;
+      strides.insert(std::pair<unsigned, int64_t>(pos, rhs.getValue()));
+      int64_t newSize = staticShape[lhs.getPosition()];
+      if (newSize == ShapedType::kDynamic || newSize == 0)
+        return failure();
path: compiler/src/iree/compiler/GlobalOptimization/ConvertStridedContractionToContraction.cpp
user: login: zjgarvey
body: Okay, I'll leave this as-is for now.
created_at: 2025-04-23T22:13:19Z
author_association: CONTRIBUTOR

PR comments:
user: login: nirvedhmeshram
created_at: 2025-04-22T22:10:46Z
body: The IR you added in https://github.com/iree-org/iree/issues/20600 LGTM, I think this is a great way to solve this problem early. Given that we will generate slow code without this optimization as this wont get recognized as GEMM otherwise, I think we should proceed with this. 
author_association: CONTRIBUTOR
user: login: zjgarvey
created_at: 2025-04-23T16:37:43Z
body: @qedawkins do you have some opinions on where this pattern should live in the compiler? It definitely needs to happen before selecting the lowering strategy, but I'm not sure what makes the most sense. 

I thought to simply add this to our "make-single-dispatch" pipeline, since I didn't want to add it more generally if it wasn't desirable there. 
author_association: CONTRIBUTOR
user: login: zjgarvey
created_at: 2025-04-24T17:15:02Z
body: Not sure why the CI froze last time. I synced with main and I'll merge after the CI is ready.
author_association: CONTRIBUTOR