AIRChannelCanonicalization: fixup a bug where canonicalization is skipped by mistake (Xilinx#1053)

erwei-xilinx · web-flow · commit 0ff42e82a9e0 · 2025-07-28T19:58:05.000-07:00
* Fixup a bug where, when num_dims &gt; max_num_dims and highest-dimension repeat is enabled, wrap-and-stride canonicalization is skipped

* Fixup test for AIE1 target: AIE1 doesn't support md dma at shim
diff --git a/mlir/include/air/Transform/AIRDependencyScheduleOpt.h b/mlir/include/air/Transform/AIRDependencyScheduleOpt.h
@@ -73,7 +73,8 @@ void populateAIRLoopIndexCanonicalizationPatterns(RewritePatternSet &patterns);
 // Populate patterns for canonicalizing offsets, sizes and strides in air
 // channel_interface operations.
 void populateAIRCanonicalizeChannelWrapAndStridePatterns(
-    RewritePatternSet &patterns, int &maxSize, bool &enableRepeatAtHighestDim);
+    RewritePatternSet &patterns, int &maxSize, int &maxNumDims,
+    bool &enableRepeatAtHighestDim);
 
 // Apply AIRSpecializeChannelWrapAndStridePattern on region.
 void applyAIRSpecializeChannelWrapAndStridePattern(
diff --git a/mlir/include/air/Util/Util.h b/mlir/include/air/Util/Util.h
@@ -166,7 +166,7 @@ int findLargestFactor(int num, int max);
 
 // Canonicalize wrap and stride lists, by removing redundant dimensions.
 LogicalResult canonicalizeWrapAndStrideList(
-    OpBuilder builder, SmallVector<Value> &offsets, SmallVector<Value> &sizes,
+    OpBuilder &builder, SmallVector<Value> &offsets, SmallVector<Value> &sizes,
     SmallVector<Value> &strides, int memref_volume, int maxSize = -1);
 
 // If wrap-and-stride lists are empty, populate them with default data access
diff --git a/mlir/lib/Conversion/AIRToAIEPass.cpp b/mlir/lib/Conversion/AIRToAIEPass.cpp
@@ -736,10 +736,11 @@ void lowerAirExecute(AIE::DeviceOp d) {
   auto ctx = d->getContext();
   RewritePatternSet patterns(ctx);
   int maxSize = isa<AIE::AIE1TargetModel>(AIE::getTargetModel(d)) ? -1 : 1023;
+  int maxNumDims = isa<AIE::AIE1TargetModel>(AIE::getTargetModel(d)) ? 1 : 4;
   patterns.insert<LowerAIRExecutePattern>(ctx);
   bool enableRepeatAtHighestDim = false;
   air::populateAIRCanonicalizeChannelWrapAndStridePatterns(
-      patterns, maxSize, enableRepeatAtHighestDim);
+      patterns, maxSize, maxNumDims, enableRepeatAtHighestDim);
   (void)applyPatternsGreedily(d, std::move(patterns));
 }
 
diff --git a/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp b/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp
@@ -2247,123 +2247,100 @@ struct AIRSpecializeChannelWrapAndStrideInAffineFor
 private:
 };
 
-struct AIRCanonicalizeChannelPutOpWrapAndStrideList
-    : public OpRewritePattern<air::ChannelPutOp> {
-  using OpRewritePattern<air::ChannelPutOp>::OpRewritePattern;
-
-  AIRCanonicalizeChannelPutOpWrapAndStrideList(MLIRContext *ctx, int &maxSize,
-                                               bool &enableRepeatAtHighestDim)
-      : OpRewritePattern(ctx), maxSize(maxSize),
+/// This pattern canonicalizes the offset/size/stride lists of `OpT` channel
+/// put/get operations.
+///
+/// **Main transformations**:
+/// 1. Detect whether a "highest-dimension repeat" pattern is active (special
+///    case where the highest dimension repeats and requires padding to
+///    `maxNumDims`).
+/// 2. Canonicalize the wrap-and-stride list by invoking
+///    `canonicalizeWrapAndStrideList()`, which normalizes
+///    offsets/sizes/strides.
+/// 3. If highest-dimension repeat is active, extend the rank of
+///    offsets/sizes/strides to `maxNumDims` by inserting zeros/ones.
+/// 4. Recreate the `OpT` operation with the canonicalized parameters and
+///    replace the original operation.
+template <typename OpT>
+struct AIRCanonicalizeChannelPutGetOpWrapAndStrideList
+    : public OpRewritePattern<OpT> {
+  using OpRewritePattern<OpT>::OpRewritePattern;
+
+  AIRCanonicalizeChannelPutGetOpWrapAndStrideList(
+      MLIRContext *ctx, int &maxSize, int &maxNumDims,
+      bool &enableRepeatAtHighestDim)
+      : OpRewritePattern<OpT>(ctx), maxSize(maxSize), maxNumDims(maxNumDims),
         enableRepeatAtHighestDim(enableRepeatAtHighestDim) {}
 
-  LogicalResult matchAndRewrite(air::ChannelPutOp op,
+  LogicalResult matchAndRewrite(OpT op,
                                 PatternRewriter &rewriter) const override {
-
+    // Collect async token types and dependencies if the op is asynchronous.
     SmallVector<Value, 1> deps;
     SmallVector<Type, 1> tys;
     if (isAsyncOp(op)) {
       tys.push_back(air::AsyncTokenType::get(op->getContext()));
       deps = op.getAsyncDependencies();
     }
 
+    // Extract offsets, sizes, and strides from the op.
     SmallVector<Value> offsets = op.getOffsets();
     SmallVector<Value> sizes = op.getSizes();
     SmallVector<Value> strides = op.getStrides();
 
-    // When highest-dimension repeat is active, (1) enableRepeatAtHighestDim
-    // option is switched on, (2) wrap-and-stride list isn't empty (i.e. data
-    // isn't 1-d streamed in), (3) highest stride is zero, and (4) highest wrap
-    // is not one.
+    // Detect if highest-dimension repeat logic should be applied.
+    // This is true when:
+    //  (1) The option enableRepeatAtHighestDim is set,
+    //  (2) The stride list is not empty,
+    //  (3) The highest (first) stride is 0, indicating repeat dimension,
+    //  (4) The highest (first) size is not 1, indicating non-zero repetition.
     bool highestDimRepeatActive = enableRepeatAtHighestDim &&
                                   !strides.empty() &&
                                   *getConstantIntValue(strides.front()) == 0 &&
                                   *getConstantIntValue(sizes.front()) != 1;
-
-    if (highestDimRepeatActive) {
-      // If repeat is enabled at the highest dimension, then the highest
-      // dimension must be preserved.
+    // If highest-dimension repeat is active but the op already has the maximum
+    // number of dimensions, no rewrite is needed.
+    if (highestDimRepeatActive && (int)offsets.size() == maxNumDims) {
       return failure();
-    }
-
-    if (failed(canonicalizeWrapAndStrideList(
-            rewriter, offsets, sizes, strides,
-            air::getTensorVolume(op.getMemref().getType()), maxSize)))
-      return failure();
-
-    auto new_op = rewriter.create<air::ChannelPutOp>(
-        op->getLoc(), tys, deps, op.getChanName(), op.getIndices(),
-        op.getMemref(), offsets, sizes, strides);
-    new_op->setAttrs(op->getDiscardableAttrDictionary());
-    for (unsigned i = 0; i < op->getResults().size(); i++)
-      op->getResults()[i].replaceAllUsesWith(new_op->getResults()[i]);
-
-    rewriter.eraseOp(op);
-
-    return success();
-  }
-
-private:
-  int &maxSize;
-  bool &enableRepeatAtHighestDim;
-};
-
-struct AIRCanonicalizeChannelGetOpWrapAndStrideList
-    : public OpRewritePattern<air::ChannelGetOp> {
-  using OpRewritePattern<air::ChannelGetOp>::OpRewritePattern;
-
-  AIRCanonicalizeChannelGetOpWrapAndStrideList(MLIRContext *ctx, int &maxSize,
-                                               bool &enableRepeatAtHighestDim)
-      : OpRewritePattern(ctx), maxSize(maxSize),
-        enableRepeatAtHighestDim(enableRepeatAtHighestDim) {}
-
-  LogicalResult matchAndRewrite(air::ChannelGetOp op,
-                                PatternRewriter &rewriter) const override {
-
-    SmallVector<Value, 1> deps;
-    SmallVector<Type, 1> tys;
-    if (isAsyncOp(op)) {
-      tys.push_back(air::AsyncTokenType::get(op->getContext()));
-      deps = op.getAsyncDependencies();
-    }
-
-    SmallVector<Value> offsets = op.getOffsets();
-    SmallVector<Value> sizes = op.getSizes();
-    SmallVector<Value> strides = op.getStrides();
-
-    // When highest-dimension repeat is active, (1) enableRepeatAtHighestDim
-    // option is switched on, (2) wrap-and-stride list isn't empty (i.e. data
-    // isn't 1-d streamed in), (3) highest stride is zero, and (4) highest wrap
-    // is not one.
-    bool highestDimRepeatActive = enableRepeatAtHighestDim &&
-                                  !strides.empty() &&
-                                  *getConstantIntValue(strides.front()) == 0 &&
-                                  *getConstantIntValue(sizes.front()) != 1;
+    } else {
+      // Canonicalize offsets/sizes/strides using a helper function.
+      if (failed(canonicalizeWrapAndStrideList(
+              rewriter, offsets, sizes, strides,
+              air::getTensorVolume(op.getMemref().getType()), maxSize)))
+        return failure();
 
-    if (highestDimRepeatActive) {
-      // If repeat is enabled at the highest dimension, then the highest
-      // dimension must be preserved.
-      return failure();
+      // When highest-dimension repeat is active, pad offsets/sizes/strides to
+      // match maxNumDims by inserting:
+      //  - offset = 0
+      //  - size   = 1
+      //  - stride = 0
+      if (highestDimRepeatActive) {
+        auto zeroIdx = rewriter.create<arith::ConstantIndexOp>(op->getLoc(), 0);
+        auto oneIdx = rewriter.create<arith::ConstantIndexOp>(op->getLoc(), 1);
+        while ((int)offsets.size() < maxNumDims) {
+          offsets.insert(offsets.begin() + 1, zeroIdx);
+        }
+        while ((int)sizes.size() < maxNumDims) {
+          sizes.insert(sizes.begin() + 1, oneIdx);
+        }
+        while ((int)strides.size() < maxNumDims) {
+          strides.insert(strides.begin() + 1, zeroIdx);
+        }
+      }
     }
 
-    if (failed(canonicalizeWrapAndStrideList(
-            rewriter, offsets, sizes, strides,
-            air::getTensorVolume(op.getMemref().getType()), maxSize)))
-      return failure();
-
-    auto new_op = rewriter.create<air::ChannelGetOp>(
-        op->getLoc(), tys, deps, op.getChanName(), op.getIndices(),
-        op.getMemref(), offsets, sizes, strides);
-    new_op->setAttrs(op->getDiscardableAttrDictionary());
-    for (unsigned i = 0; i < op->getResults().size(); i++)
-      op->getResults()[i].replaceAllUsesWith(new_op->getResults()[i]);
-
-    rewriter.eraseOp(op);
+    // Create a new op with the canonicalized attributes and operands.
+    auto attrs = op->getDiscardableAttrDictionary();
+    auto new_op = rewriter.replaceOpWithNewOp<OpT>(
+        op, tys, deps, op.getChanName(), op.getIndices(), op.getMemref(),
+        offsets, sizes, strides);
+    new_op->setAttrs(attrs);
 
     return success();
   }
 
 private:
   int &maxSize;
+  int &maxNumDims;
   bool &enableRepeatAtHighestDim;
 };
 
@@ -3171,7 +3148,7 @@ LogicalResult AIRSpecializeChannelWrapAndStrideImpl(
   // Canonicalize wrap and stride list to remove redundant dimensions
   RewritePatternSet preproc_wns_patterns(ctx);
   populateAIRCanonicalizeChannelWrapAndStridePatterns(
-      preproc_wns_patterns, maxSize, enableRepeatAtHighestDim);
+      preproc_wns_patterns, maxSize, maxNumDims, enableRepeatAtHighestDim);
   (void)applyPatternsGreedily(*region, std::move(preproc_wns_patterns));
 
   RewritePatternSet patterns(ctx);
@@ -3196,10 +3173,9 @@ LogicalResult AIRSpecializeChannelWrapAndStrideImpl(
 
   // Canonicalize wrap and stride list to remove redundant dimensions
   RewritePatternSet cano_patterns(ctx);
-  populateAIRCanonicalizeChannelWrapAndStridePatterns(cano_patterns, maxSize,
-                                                      enableRepeatAtHighestDim);
+  populateAIRCanonicalizeChannelWrapAndStridePatterns(
+      cano_patterns, maxSize, maxNumDims, enableRepeatAtHighestDim);
   ExecuteOp::getCanonicalizationPatterns(cano_patterns, ctx);
-  // WaitAllOp::getCanonicalizationPatterns(cano_patterns, ctx);
   (void)applyPatternsGreedily(*region, std::move(cano_patterns));
 
   return success();
@@ -6733,11 +6709,13 @@ void populateAIRLoopIndexCanonicalizationPatterns(RewritePatternSet &patterns) {
 }
 
 void populateAIRCanonicalizeChannelWrapAndStridePatterns(
-    RewritePatternSet &patterns, int &maxSize, bool &enableRepeatAtHighestDim) {
+    RewritePatternSet &patterns, int &maxSize, int &maxNumDims,
+    bool &enableRepeatAtHighestDim) {
   MLIRContext *ctx = patterns.getContext();
-  patterns.insert<AIRCanonicalizeChannelPutOpWrapAndStrideList,
-                  AIRCanonicalizeChannelGetOpWrapAndStrideList>(
-      ctx, maxSize, enableRepeatAtHighestDim);
+  patterns.insert<
+      AIRCanonicalizeChannelPutGetOpWrapAndStrideList<air::ChannelPutOp>,
+      AIRCanonicalizeChannelPutGetOpWrapAndStrideList<air::ChannelGetOp>>(
+      ctx, maxSize, maxNumDims, enableRepeatAtHighestDim);
 }
 
 void applyAIRSpecializeChannelWrapAndStridePattern(
diff --git a/mlir/lib/Util/Util.cpp b/mlir/lib/Util/Util.cpp
@@ -1086,11 +1086,12 @@ int air::findLargestFactor(int num, int max) {
 
 // Canonicalize wrap and stride lists by removing redundant dimensions.
 LogicalResult air::canonicalizeWrapAndStrideList(
-    OpBuilder builder, SmallVector<Value> &offsets, SmallVector<Value> &sizes,
+    OpBuilder &builder, SmallVector<Value> &offsets, SmallVector<Value> &sizes,
     SmallVector<Value> &strides, int memref_volume, int maxSize) {
   // AIE2 hardware constraints. TODO: import these info from target model.
   const int AIE2_STRIDE_UPPER_BOUND = 1048576;
   bool listsHaveChanged = false;
+  OpBuilder::InsertionGuard guard(builder);
   // Match offsets size with sizes and strides
   auto max_dim_size =
       std::max(std::max(offsets.size(), sizes.size()), strides.size());
diff --git a/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds.mlir b/mlir/test/Transform/AIRDependencyScheduleOpt/opt_shim_dma_bds.mlir
@@ -855,4 +855,29 @@ module {
     }
     return
   }
+
+  // Canonicalizing repeat dimension at highest dimension.
+
+  // CHECK-LABEL: func15
+  // CHECK: air.channel.put async{{.*}}@channel_0[%c0{{.*}}, %c0{{.*}}] (%{{.*}}[%c0{{.*}}, %c0{{.*}}, %c0{{.*}}, %c320{{.*}}] [%c2{{.*}}, %c1{{.*}}, %c512{{.*}}, %c64{{.*}}] [%c0{{.*}}, %c0{{.*}}, %c512{{.*}}, %c1{{.*}}])
+  // NPUTILED-LABEL: func15
+  // NPUTILED: air.channel.put async{{.*}}@channel_0[%c0{{.*}}, %c0{{.*}}] (%{{.*}}[%c0{{.*}}, %c0{{.*}}, %c0{{.*}}, %c320{{.*}}] [%c2{{.*}}, %c1{{.*}}, %c512{{.*}}, %c64{{.*}}] [%c0{{.*}}, %c0{{.*}}, %c512{{.*}}, %c1{{.*}}])
+  // AIE1-LABEL: func15
+  // AIE1: air.channel.put async{{.*}}@channel_0[%c0{{.*}}, %c0{{.*}}] (%{{.*}}[%c0{{.*}}, %c0{{.*}}, %c320{{.*}}] [%c2{{.*}}, %c512{{.*}}, %c64{{.*}}] [%c0{{.*}}, %c512{{.*}}, %c1{{.*}}])
+
+  func.func @func15(%arg0: memref<512x512xbf16>) {
+    %0 = air.launch async () in () args(%arg8=%arg0) : memref<512x512xbf16> {
+      %c65536 = arith.constant 65536 : index
+      %c4 = arith.constant 4 : index
+      %c256 = arith.constant 256 : index
+      %c64 = arith.constant 64 : index
+      %c128 = arith.constant 128 : index
+      %c512 = arith.constant 512 : index
+      %c0 = arith.constant 0 : index
+      %c1 = arith.constant 1 : index
+      %c2 = arith.constant 2 : index
+      %1 = air.channel.put async  @channel_0[%c0, %c0] (%arg8[%c0, %c0, %c1, %c0, %c256] [%c2, %c4, %c1, %c128, %c64] [%c0, %c65536, %c64, %c512, %c1]) {id = 6 : i32, metadataArray = [{base = "air_channel_13_0", index = 0 : i32}, {base = "air_channel_13_1", index = 1 : i32}, {base = "air_channel_13_2", index = 2 : i32}, {base = "air_channel_13_3", index = 3 : i32}]} : (memref<512x512xbf16>)
+    }
+    return
+  }
 }
