[Mlir-commits] [mlir] 5198205 - [mlir][affine] Make [de]linearize_index a valid source of dims (#138929)

[llvmlistbot at llvm.org]

Author: Krzysztof Drewniak
Date: 2025-05-13T11:13:08-05:00
New Revision: 5198205fb7d0228fdab2c08a24d22e97de673e51

URL: https://github.com/llvm/llvm-project/commit/5198205fb7d0228fdab2c08a24d22e97de673e51
DIFF: https://github.com/llvm/llvm-project/commit/5198205fb7d0228fdab2c08a24d22e97de673e51.diff

LOG: [mlir][affine] Make [de]linearize_index a valid source of dims (#138929)

There's a sense in which affine.linearize_index and
affine.delinearize_index are special-cases of affine.apply (which get
their own ops to enable better code generation and more accurate
canonicalization). Therefore, allow these operations to be dimension
operands for operations like affine.load just like affine.apply can be.

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/invalid.mlir
    mlir/test/Dialect/Affine/ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 6f99306b54b45..9cb06e014f843 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -305,6 +305,8 @@ bool mlir::affine::isValidDim(Value value) {
 // *) It is valid as a symbol.
 // *) It is an induction variable.
 // *) It is the result of an affine apply operation with dimension id operands.
+// *) It is the result of a more specialized index transformation (ex.
+// delinearize_index or linearize_index) with dimension id operands.
 bool mlir::affine::isValidDim(Value value, Region *region) {
   // The value must be an index type.
   if (!value.getType().isIndex())
@@ -325,6 +327,11 @@ bool mlir::affine::isValidDim(Value value, Region *region) {
   // Affine apply operation is ok if all of its operands are ok.
   if (auto applyOp = dyn_cast<AffineApplyOp>(op))
     return applyOp.isValidDim(region);
+  // delinearize_index and linearize_index are special forms of apply
+  // and so are valid dimensions if all their arguments are valid dimensions.
+  if (isa<AffineDelinearizeIndexOp, AffineLinearizeIndexOp>(op))
+    return llvm::all_of(op->getOperands(),
+                        [&](Value arg) { return ::isValidDim(arg, region); });
   // The dim op is okay if its operand memref/tensor is defined at the top
   // level.
   if (auto dimOp = dyn_cast<ShapedDimOpInterface>(op))

diff  --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index 9703c05fff8f6..c4bb22f9a8dae 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -544,6 +544,34 @@ func.func @dynamic_dimension_index() {
 
 // -----
 
+func.func @dynamic_linearized_index() {
+  "unknown.region"() ({
+    %idx = "unknown.test"() : () -> (index)
+    %memref = "unknown.test"() : () -> memref<?xf32>
+    %pos = affine.linearize_index [%idx, %idx] by (8) : index
+    // expected-error at below {{op operand cannot be used as a dimension id}}
+    affine.load %memref[%pos] : memref<?xf32>
+    "unknown.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+func.func @dynamic_delinearized_index() {
+  "unknown.region"() ({
+    %idx = "unknown.test"() : () -> (index)
+    %memref = "unknown.test"() : () -> memref<?x?xf32>
+    %pos0, %pos1 = affine.delinearize_index %idx into (8) : index, index
+    // expected-error at below {{op operand cannot be used as a dimension id}}
+    affine.load %memref[%pos0, %pos1] : memref<?x?xf32>
+    "unknown.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
 #map = affine_map<() -> ()>
 #map1 = affine_map<() -> (1)>
 func.func @no_lower_bound() {

diff  --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir
index 233c18c82831c..8a3f41d1d9b05 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -148,6 +148,23 @@ func.func @valid_symbol_affine_scope(%n : index, %A : memref<?xf32>) {
 
 // -----
 
+// Test dimension constraints for linearize_index and delinearize_index
+
+// CHECK-LABEL: func @valid_dim_linearize_delinearize
+func.func @valid_dim_linearize_delinearize(%m : index, %n : index, %A : memref<?xf32>, %B: memref<?x32x?xf32>) {
+    affine.for %0 = 0 to %m {
+      affine.for %1 = 0 to %n {
+        %load_idx = affine.linearize_index disjoint [%0, %1] by (%m, %n) : index
+        %store_idx0, %store_idx1 = affine.delinearize_index %n into (32) : index, index
+        %v = affine.load %A[%load_idx] : memref<?xf32>
+        affine.store %v, %B[%0, %store_idx1, %store_idx0] : memref<?x32x?xf32>
+      }
+    }
+  return
+}
+
+// -----
+
 // Test the fact that module op always provides an affine scope.
 
 %idx = "test.foo"() : () -> (index)
@@ -309,7 +326,7 @@ func.func @linearize_mixed(%index0: index, %index1: index, %index2: index, %basi
 module {
   func.func @gpu_launch_affine() {
     %c1 = arith.constant 1 : index
-    gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) 
+    gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
     threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
       %thread_id_x = gpu.thread_id  x
       %c128 = arith.constant 128 : index