[mlir][affine] Make [de]linearize_index a valid source of dims

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -307,6 +307,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())
@@ -326,6 +328,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))

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@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@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() {

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