support 2d specialization (Xilinx#500)

nod-ai · Mar 18, 2024 · c73fd4b · c73fd4b
1 parent 22306b3
commit c73fd4b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 20 deletions.
diff --git a/mlir/lib/Conversion/AIRToAIEPass.cpp b/mlir/lib/Conversion/AIRToAIEPass.cpp
@@ -1804,22 +1804,20 @@ class AIRToAIEPass : public air::impl::AIRToAIEBase<AIRToAIEPass> {
                         dyn_cast<mlir::affine::AffineApplyOp>(child_op))
                   position_apply = apply_op;
               }
-              int position_iv = -1;
-              if (position_apply)
-                for (unsigned i = 0; i < par.getInductionVars().size(); i++)
-                  for (auto map_o : position_apply.getMapOperands())
-                    if (par.getInductionVars()[i] == map_o)
-                      position_iv = i;
-              if (position_apply && position_iv >= 0) {
-                // Evaluate the affine expression and specialize the operand
-                auto c = position_apply.getAffineMap().getResult(0);
-                SmallVector<AffineExpr, 1> const_syms{
-                    getAffineConstantExpr(position[position_iv],
-                                          builder.getContext()),
-                };
-                auto newC = c.replaceSymbols(const_syms);
-                auto expr = dyn_cast<AffineConstantExpr>(
-                    simplifyAffineExpr(newC, 0, 1));
+              if (position_apply) {
+                SmallVector<AffineExpr> const_syms;
+                for (unsigned i = 0; i < par.getInductionVars().size(); i++) {
+                  for (auto map_o : position_apply.getMapOperands()) {
+                    if (par.getInductionVars()[i] == map_o) {
+                      const_syms.push_back(getAffineConstantExpr(
+                          position[i], builder.getContext()));
+                    }
+                  }
+                }
+                AffineExpr newC = position_apply.getAffineMap().getResult(0);
+                newC = newC.replaceSymbols(const_syms);
+                auto expr = dyn_cast<AffineConstantExpr>(simplifyAffineExpr(
+                    newC, 0, position_apply.getMapOperands().size()));
                 assert(expr);
                 int result = expr.getValue();
                 remap.map(oper, builder.create<arith::ConstantIndexOp>(

diff --git a/test/airhost/40_air_8x4_2d_square/air.mlir b/test/airhost/40_air_8x4_2d_square/air.mlir
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#map0 = affine_map<(d0,d1) -> (d0 * 32 + d1 * 8)>
+#map0 = affine_map<()[d0,d1] -> (d0 * 32 + d1 * 8)>
 module {
 
 func.func @graph(%arg0 : memref<256x16xi32>, %arg1 : memref<256x16xi32>) -> () {
@@ -19,7 +19,7 @@ func.func @graph(%arg0 : memref<256x16xi32>, %arg1 : memref<256x16xi32>) -> () {
     %c8 = arith.constant 8 : index
     %buf0 = memref.alloc() {sym_name = "scratch"}: memref<8x8xi32, 2>
     %buf1 = memref.alloc() {sym_name = "scratch_copy"}: memref<8x8xi32, 2>
-    %5 = affine.apply #map0(%tx, %ty)
+    %5 = affine.apply #map0()[%tx, %ty]
     air.dma_memcpy_nd (%buf0[%c0, %c0][%c8, %c8][%c8, %c0], %ext0[%c0, %5][%c8, %c8][%c256, %c0]) {id = 1 : i32} : (memref<8x8xi32, 2>, memref<256x16xi32>)
     affine.for %j = 0 to 8 {
       affine.for %i = 0 to 8 {
@@ -36,5 +36,3 @@ func.func @graph(%arg0 : memref<256x16xi32>, %arg1 : memref<256x16xi32>) -> () {
 }
 
 }
-
-