only rlvnt

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIETileAndFuse.cpp

@@ -16,12 +16,12 @@
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
-#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Iterators.h"
 #include "mlir/IR/PatternMatch.h"
 
 #define DEBUG_TYPE "iree-amdaie-tile-and-fuse"
 
+
 namespace mlir::iree_compiler::AMDAIE {
 
 namespace {
@@ -157,43 +157,27 @@ void AMDAIETileAndFusePass::runOnOperation() {
 
   SmallVector<OpFoldResult> tileSizes =
       getAsIndexOpFoldResult(context, tileSizesVal);
-
   auto options = scf::SCFTilingOptions().setTileSizes(tileSizes);
 
   // When tiling using scf.for we do not need to set any mapping.
   if (!useSCFFor) {
     options.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);
-
-    // There is one induction variables in the scf.forall for each of the
-    // non-zero tile sizes. Recall that a '0' tile size corresonds to 'do
-    // not tile'.
-    uint32_t nbIndVars = std::count_if(tileSizesVal.begin(), tileSizesVal.end(),
-                                       [](int64_t t) { return t != 0; });
-
-    // See mlir::gpu::MappingId enum: there are currently 13 values.
-    if (nbIndVars > mlir::gpu::getMaxEnumValForMappingId()) {
-      LLVM_DEBUG(llvm::dbgs() << "----- skip, too many loops to tile -----\n");
-      return;
-    }
-    auto getMappingAttributeForDimension = [&](uint32_t i) -> Attribute {
-      auto id = static_cast<gpu::MappingId>(i);
-      if (tilingLevel == 0) return gpu::GPUBlockMappingAttr::get(context, id);
-      return gpu::GPUThreadMappingAttr::get(context, id);
-    };
-
-    SmallVector<Attribute> mapping;
-    mapping.reserve(nbIndVars);
-    for (uint32_t i = 0; i < nbIndVars; ++i) {
-      // DimY and DimX are swapped, so that DimX goes to AIE array columns and
-      // DimX goes to AIE arrat rows (or something like that...)
-      if (i == 0)
-        mapping.push_back(getMappingAttributeForDimension(1));
-      else if (i == 1)
-        mapping.push_back(getMappingAttributeForDimension(0));
-      else
-        mapping.push_back(getMappingAttributeForDimension(i));
+    // Here we assume there are always two levels of parallel (scf.forall)
+    // loops, and the first level of tiling is always using scf.forall and
+    // mapped to blocks. Currently we are not using mapping attributes for
+    // Conv2d ops, because there could be four parallel tiling dimensions.
+    // TODO (vivian): create AIE specific mapping attributes.
+    if (!isa<linalg::ConvolutionOpInterface>(consumerOp.getOperation())) {
+      if (tilingLevel == 0) {
+        options.setMapping(
+            {gpu::GPUBlockMappingAttr::get(context, gpu::MappingId::DimY),
+             gpu::GPUBlockMappingAttr::get(context, gpu::MappingId::DimX)});
+      } else {
+        options.setMapping(
+            {gpu::GPUThreadMappingAttr::get(context, gpu::MappingId::DimY),
+             gpu::GPUThreadMappingAttr::get(context, gpu::MappingId::DimX)});
+      }
     }
-    options.setMapping(mapping);
   }
 
   IRRewriter rewriter(context);