Add support of param type for transform.structured.tile_using_forall #72097

jinchen62 · 2023-11-13T10:29:56Z

Make transform.structured.tile_using_forall be able to take param type tile sizes.

Examples:

%tile_sizes = transform.param.constant 16 : i64 -> !transform.param<i64>
transform.structured.tile_using_forall %matmul tile_sizes [%tile_sizes : !transform.param<i64>, 32] ( mapping = [#gpu.block<x>, #gpu.block<y>] ) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)

%c10 = transform.param.constant 10 : i64 -> !transform.any_param
%c20 = transform.param.constant 20 : i64 -> !transform.any_param
%tile_sizes = transform.merge_handles %c10, %c20 : !transform.any_param
transform.structured.tile_using_forall %matmul tile_sizes *(%tile_sizes : !transform.any_param) ( mapping = [#gpu.block<x>, #gpu.block<y>] ) : (!transform.any_op) -> (!transform.any_op, !transform.any_op)

llvmbot · 2023-11-13T10:30:25Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-linalg

Author: jinchen62 (jinchen62)

Changes

Make transform.structured.tile_using_forall be able to take param type tile sizes.

Examples:

%tile_size1 = transform.param.constant 16 : i64 -&gt; !transform.param&lt;i64&gt;
transform.structured.tile_using_forall %matmul tile_sizes [%tile_size1 : !transform.param&lt;i64&gt;, 32] ( mapping = [#gpu.block&lt;x&gt;, #gpu.block&lt;y&gt;] ) : (!transform.any_op) -&gt; (!transform.any_op, !transform.any_op)

%tile_sizes = transform.param.constant [16 : i64, 32 : i64] -&gt; !transform.any_param
transform.structured.tile_using_forall %matmul tile_sizes [%tile_sizes : !transform.any_param] ( mapping = [#gpu.block&lt;x&gt;, #gpu.block&lt;y&gt;] ) : (!transform.any_op) -&gt; (!transform.any_op, !transform.any_op)

Full diff: https://github.com/llvm/llvm-project/pull/72097.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+5-5)
(modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+24-2)

diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index f1c3d717f1fa951..a24f6ff8308ba34 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -23,7 +23,7 @@ include "mlir/IR/RegionKindInterface.td"
 // value in the payload IR.
 def TransformParamTypeOrAnyHandle : Type<
     Or<[TransformHandleTypeInterface.predicate,
-        Transform_ParamType.predicate]>,
+        TransformParamTypeInterface.predicate]>,
     "transform 'param' type or any handle type">;
 
 //===----------------------------------------------------------------------===//
@@ -1924,10 +1924,10 @@ def TileUsingForallOp :
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
-                   Variadic<TransformHandleTypeInterface>:$num_threads,
-                   Variadic<TransformHandleTypeInterface>:$tile_sizes,
-                   Optional<TransformHandleTypeInterface>:$packed_num_threads,
-                   Optional<TransformHandleTypeInterface>:$packed_tile_sizes,
+                   Variadic<TransformParamTypeOrAnyHandle>:$num_threads,
+                   Variadic<TransformParamTypeOrAnyHandle>:$tile_sizes,
+                   Optional<TransformParamTypeOrAnyHandle>:$packed_num_threads,
+                   Optional<TransformParamTypeOrAnyHandle>:$packed_tile_sizes,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_num_threads,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_tile_sizes,
                    OptionalAttr<DeviceMappingArrayAttr>:$mapping);
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index de4965f937162ea..4bf4db3381fab79 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -98,12 +98,34 @@ static DiagnosedSilenceableFailure unpackSingleIndexResultPayloadOperations(
       result.push_back(ofr);
       continue;
     }
-    auto payloadOps = state.getPayloadOps(ofr.get<Value>());
+
+    Value transformValue = ofr.get<Value>();
+    if (isa<ParamType>(transformValue.getType())) {
+      ArrayRef<Attribute> params = state.getParams(transformValue);
+      if (!isa<IntegerAttr>(params[0]))
+        return transformOp.emitDefiniteFailure() << "expected IntegerAttr";
+      result.push_back(params[0]);
+      continue;
+    }
+    if (isa<AnyParamType>(transformValue.getType())) {
+      ArrayRef<Attribute> params = state.getParams(transformValue);
+      if (!isa<ArrayAttr>(params[0]))
+        return transformOp.emitDefiniteFailure() << "expected ArrayAttr";
+      ArrayAttr paramsArray = cast<ArrayAttr>(params[0]);
+      for (Attribute param : paramsArray.getValue()) {
+        if (!isa<IntegerAttr>(param))
+          return transformOp.emitDefiniteFailure() << "expected IntegerAttr";
+        result.push_back(param);
+      }
+      continue;
+    }
+
+    auto payloadOps = state.getPayloadOps(transformValue);
     if (!llvm::hasSingleElement(payloadOps)) {
       DiagnosedSilenceableFailure diag =
           transformOp.emitSilenceableError()
           << "handle must be mapped to exactly one payload op";
-      diag.attachNote(ofr.get<Value>().getLoc())
+      diag.attachNote(transformValue.getLoc())
           << "mapped to " << llvm::range_size(payloadOps) << " payload ops";
       return diag;
     }

ftynse

Thanks for you PR! Please address inline comments.

Also please add tests. The good practice is to have at least one non-trivial test for desired behavior, tests for edge cases, and tests for all user-visible diagnostic messages.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

ftynse · 2023-11-13T12:12:23Z

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

+      for (Attribute param : paramsArray.getValue()) {
+        if (!isa<IntegerAttr>(param))
+          return transformOp.emitDefiniteFailure() << "expected IntegerAttr";
+        result.push_back(param);


It's unclear to me what is the intended model here. Looking at the code, if the type is any_param, it will accept array-of-integer parameters and expand its context as tile sizes. But it will keep doing so for other parameters, so we may end up having more sizes than loops, and a surprising mismatch between parameter positions and loop depth. We could accept one parameter, and unpack it, but I don't think we want to accept multiple here. Also, there's no need to rely on array attributes as a parameter can be associated with a list of attributes (unless we want to support some multi-level scheme where the parameter is a list-of-lists of tile sizes to use for each of the loops associated with the input).

I don't believe this was addressed.

I thought moving the support of any_param to the second unpackSingleIndexResultPayloadOperations that takes a Value would address it since there won't be lists inside lists. Could you give me more guidence on this?

As far as I recall, there are currently two modes of functioning here:

there are N unpacked handles, where each handle corresponds to one dimension; each handle points to the same number of operations as the "target" handle;

there is 1 packed handle, which points to as many operations as dimensions; same dimensions are used for tiling all operations associated with the "target" handle.

Parameters should behave similarly. Since parameters are associated with a list of attributes, we would treat that list identically to how we treat the list of operations for operation handle. Consequently, we should never expect an array attribute in parameters. There is already the implicit list and we must use that instead.

Please review, thanks!

ftynse · 2023-11-13T12:13:58Z

Also note that the utility function that is changed in this PR is used in more cases than TileUsingForall, make sure to test those as well.

github-actions · 2023-11-15T22:53:03Z

✅ With the latest revision this PR passed the C/C++ code formatter.

jinchen62 · 2023-11-16T22:30:26Z

@ftynse Thanks for the review! I've addressed the comments, please check.

ftynse

Please add tests for user-visible error messages (maybe you forgot to git add a file with those?)

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

ftynse

Thank you! It's almost ready to go, please address the remaining comments.

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

mlir/test/Dialect/Linalg/tile-to-forall.mlir

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

jinchen62 · 2024-01-25T18:02:45Z

@ftynse Please review, thanks!

ftynse

Thanks for iterating on this!

ftynse · 2024-01-31T09:02:36Z

Windows problems are due to a missing numpy installation and are irrelevant here.

llvmbot added mlir:linalg mlir labels Nov 13, 2023

ftynse requested changes Nov 13, 2023

View reviewed changes

jinchen62 force-pushed the param branch from 5f51dfe to d53f53c Compare November 15, 2023 22:50

jinchen62 force-pushed the param branch 3 times, most recently from a04b3bd to 8d09cc2 Compare November 16, 2023 22:28

jinchen62 requested a review from ftynse November 16, 2023 22:29

jinchen62 force-pushed the param branch 2 times, most recently from e50f7c3 to 48c045e Compare November 17, 2023 05:12

ftynse requested changes Nov 17, 2023

View reviewed changes

jinchen62 force-pushed the param branch 2 times, most recently from c4be997 to 3a524ab Compare November 30, 2023 19:56

jinchen62 force-pushed the param branch from 3a524ab to ef15dae Compare January 9, 2024 15:51

jinchen62 requested review from dcaballe and nicolasvasilache as code owners January 9, 2024 15:51

jinchen62 force-pushed the param branch 2 times, most recently from 27a4e49 to b0d4192 Compare January 17, 2024 23:26

jinchen62 requested a review from ftynse January 17, 2024 23:29

jinchen62 force-pushed the param branch from b0d4192 to 831c728 Compare January 19, 2024 00:02

ftynse reviewed Jan 19, 2024

View reviewed changes

jinchen62 force-pushed the param branch 4 times, most recently from fa63fb8 to 4851368 Compare January 24, 2024 10:17

jinchen62 requested a review from ftynse January 24, 2024 10:18

jinchen62 force-pushed the param branch 3 times, most recently from 988b39e to dcb10ea Compare January 25, 2024 18:02

jinchen62 force-pushed the param branch from dcb10ea to df347ad Compare January 29, 2024 02:30

ftynse approved these changes Jan 30, 2024

View reviewed changes

Add support of param type for transform.structured.tile_using_forall

64e2acb

ftynse force-pushed the param branch from df347ad to 64e2acb Compare January 30, 2024 16:04

ftynse merged commit d439f36 into llvm:main Jan 31, 2024

jinchen62 deleted the param branch January 31, 2024 21:18

Add support of param type for transform.structured.tile_using_forall #72097

Add support of param type for transform.structured.tile_using_forall #72097

Uh oh!

Conversation

jinchen62 commented Nov 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Nov 13, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ftynse Nov 13, 2023

Choose a reason for hiding this comment

Uh oh!

ftynse Nov 17, 2023

Choose a reason for hiding this comment

Uh oh!

jinchen62 Nov 29, 2023

Choose a reason for hiding this comment

Uh oh!

ftynse Dec 1, 2023

Choose a reason for hiding this comment

Uh oh!

jinchen62 Jan 17, 2024

Choose a reason for hiding this comment

Uh oh!

ftynse commented Nov 13, 2023

Uh oh!

github-actions bot commented Nov 15, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jinchen62 commented Nov 16, 2023

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jinchen62 commented Jan 25, 2024

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

ftynse commented Jan 31, 2024

Uh oh!

Uh oh!

jinchen62 commented Nov 13, 2023 •

edited

Loading

llvmbot commented Nov 13, 2023 •

edited

Loading

github-actions bot commented Nov 15, 2023 •

edited

Loading