Fixes for llvm/llvm-project#143467

Signed-off-by: MaheshRavishankar <[email protected]>

Author: MaheshRavishankar

compiler/src/iree/compiler/Codegen/Common/GPU/GPUApplyTilingLevel.cpp

@@ -132,8 +132,15 @@ static LogicalResult applyTileAndFuseToEachRoot(
 
     if (tilingLevel == IREE::GPU::TilingLevel::PartialReduction) {
       tilingOptions.setReductionTilingStrategy(
-          scf::SCFTilingOptions::ReductionTilingStrategy::
-              PartialReductionOuterReduction);
+          ReductionTilingStrategy::PartialReductionOuterReduction);
+      SmallVector<unsigned> reductionDims;
+      for (auto [index, iteratorType] :
+           llvm::enumerate(tilingInterfaceOp.getLoopIteratorTypes())) {
+        if (iteratorType == utils::IteratorType::reduction) {
+          reductionDims.push_back(index);
+        }
+      }
+      tilingOptions.setReductionDims(reductionDims);
     }
 
     scf::SCFTileAndFuseOptions tileAndFuseOptions;

compiler/src/iree/compiler/Dialect/LinalgExt/IR/TilingInterfaceImpl.cpp

@@ -2501,7 +2501,7 @@ static AffineMap getPartialResultMap(AffineMap map, AttentionOpDetail &opInfo) {
 FailureOr<SmallVector<Value>>
 OnlineAttentionOp::generateInitialTensorForPartialReduction(
     OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
-    ArrayRef<int> reductionDim) {
+    const llvm::SetVector<unsigned> &reductionDims) {
   FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
       getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
   if (failed(maybeOpInfo)) {
@@ -2555,8 +2555,10 @@ OnlineAttentionOp::generateInitialTensorForPartialReduction(
 }
 
 FailureOr<TilingResult> OnlineAttentionOp::tileToPartialReduction(
-    OpBuilder &b, Location loc, ValueRange init, ArrayRef<OpFoldResult> offsets,
-    ArrayRef<OpFoldResult> sizes, ArrayRef<int> reductionDims) {
+    OpBuilder &b, Location loc, ReductionTilingStrategy strategy,
+    ValueRange init, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes,
+    const llvm::SetVector<unsigned> &reductionDims) {
   FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
       getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
   if (failed(maybeOpInfo)) {
@@ -2707,10 +2709,9 @@ static Value computeSubAndExp2(OpBuilder &builder, Location loc,
   return genericOp.getResult(0);
 }
 
-FailureOr<MergeResult>
-OnlineAttentionOp::mergeReductions(OpBuilder &b, Location loc,
-                                   ValueRange partialReduce,
-                                   ArrayRef<int> reductionDim) {
+FailureOr<MergeResult> OnlineAttentionOp::mergeReductions(
+    OpBuilder &b, Location loc, ValueRange partialReduce,
+    const llvm::SetVector<unsigned> &reductionDims) {
   FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
       getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());
   if (failed(maybeOpInfo)) {
@@ -2753,8 +2754,10 @@ OnlineAttentionOp::mergeReductions(OpBuilder &b, Location loc,
 
 LogicalResult OnlineAttentionOp::getPartialResultTilePosition(
     OpBuilder &b, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
-    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
-    SmallVector<OpFoldResult> &resultSizes, ArrayRef<int> reductionDims) {
+    ArrayRef<OpFoldResult> sizes,
+    const llvm::SetVector<unsigned> &reductionDims,
+    SmallVector<OpFoldResult> &resultOffsets,
+    SmallVector<OpFoldResult> &resultSizes) {
 
   FailureOr<AttentionOpDetail> maybeOpInfo = AttentionOpDetail::get(
       getQueryMap(), getKeyMap(), getValueMap(), getOutputMap());