[Question] Conditional OneToN conversion using `OneToNOpConversionPattern`

cs9614 · November 6, 2024, 4:00am

Hi, I have a case where I need to decompose a subset of ops (belong to same op type) into multiple ops that work on smaller size inputs.

For example,

%0 = create_tile ... {decompose} : tile<32x16xf16>
%1 = create_tile ... : tile<32x16xf16>
%2 = load_tile %0 : tile<32x16xf16> -> vector<32x16xf16>
%3 = load_tile %1 : tile<32x16xf16> -> vector<32x16xf16>

Only create_tile ops with {decompose} attribute must be 1:N converted. like so,

%0 = create_tile ... : tile<16x16xf16>
%1 = create_tile ...  : tile<16x16xf16>

%2 = create_tile ... : tile<32x16xf16>

%3 = load_tile %0 : tile<16x16xf16> -> vector<16x16xf16>
%4 = load_tile %1 : tile<16x16xf16> -> vector<16x16xf16>

%5 = load_tile %2 : tile<32x16xf16> -> vector<32x16xf16>

Is this doable using the OneToNOpConversionPattern interface? I tried for some time without any success. When I define the type converter to convert tile<32x16xf16> -> 2 * tile<16x16xf16>, materialization hooks end up inserting unrealized_cast ops for the ops that I need unchanged.

Resulting in something like,

%0 = create_tile ... : tile<16x16xf16>
%1 = create_tile ...  : tile<16x16xf16>

%2 = create_tile ... : tile<32x16xf16>
%3:2 = unrealized_cast %2 : tile<32x16xf16> -> tile<16x16xf16>, tile<16x16xf16>

%4 = load_tile %0 : tile<16x16xf16> -> vector<16x16xf16>
%5 = load_tile %1 : tile<16x16xf16> -> vector<16x16xf16>

%6 = load_tile %3:0 : tile<32x16xf16> -> vector<32x16xf16>
%7 = load_tile %3:1 : tile<32x16xf16> -> vector<32x16xf16>

I also tried using usual OpRewritePattern and somehow maintaining the type legality by using manual unrealized_cast insertion which works fine. But it would be much nicer if I can do this using OneToNConversion.

Any suggestions are highly appreciated.

matthias-springer · November 6, 2024, 5:32am

It looks like you need a context-specific type converter: A type converter that is only applied to some ops. Neither the 1:N nor the regular dialect conversion support that at the moment.

You could try to run the 1:N dialect conversion without a type converter (pass a nullptr type converter to the pattern constructor). That should work with the regular dialect conversion and hopefully also the 1:N dialect conversion that you need in your case.

fyi, I am in the process of merging the two dialect conversion drivers, so this part of the API is going to change slightly soon.

cs9614 · November 6, 2024, 3:48pm

Thanks for the reply.

I tried this approach, but it did not work out. Then OneToNConversion complains that in some cases there is no 1:N conversions and you have just 1:1 conversion (match in number of outputs). Looks like OneToNConversion is not designed to handle this case.

matthias-springer · November 6, 2024, 11:34pm

Did it trigger an assertion?

cs9614 · November 7, 2024, 3:51pm

yes. If my oneToNTypeConverter does not have any 1:N type conversions added. It gives me following error,.

`newValues.size() == resultMapping.getConvertedTypes().size()' failed.

My type converter is as follows,

OneToNTypeConverter typeConverter;
typeConverter.addConversion([](Type type) { return type; });
auto addNToOneCast = [](OpBuilder &builder, Type type, ValueRange inputs,
                            Location loc) {
      auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
      return std::optional<Value>(cast.getResult(0));
    };
auto addOneToNCast = [](OpBuilder &builder, TypeRange types, Value input,
                            Location loc) {
      auto cast = builder.create<UnrealizedConversionCastOp>(loc, types, input);
      return std::optional<ValueRange>(cast.getResults());
    };
typeConverter.addSourceMaterialization(addNToOneCast);
typeConverter.addArgumentMaterialization(addNToOneCast);
typeConverter.addTargetMaterialization(addOneToNCast);

If I remove,

typeConverter.addConversion([](Type type) { return type; });

my conversion pattern does not get called. Also I would like to keep the materialization hooks, because that makes handling control-flow ops much easier with signature conversions.

P.S.
I managed to get whole thing working by just using the 1:1 OpConversionPatterns, few drawbacks there are,

Need to rely on unrealized_cast to reconcile types.
Bit of a pain to handle ops that require signature conversions.
I had to define a conversion target also where I marked the subset of ops that needs to be decomposed as illegal. Can not reply on OpRewritePatterns because I needed a type converter to do the signature conversion on For Ops.

I believe to handle this in 1:N conversion, I need a custom 1:N type converter that allows to look at some of the properties of the op before converting some type. Is there any plan on supporting something like that?

matthias-springer · November 9, 2024, 10:39am

You are manually calling OneToNPatternRewriter::replaceOp, right? That’s where the assertion is failing. It looks like you are passing an incorrect resultMapping.

Actually, I don’t see how this is related to the type converter. (Or running a conversion pattern without a type converter.) OneToNPatternRewriter::replaceOp does not require a type converter.

charitha22 · January 21, 2025, 10:50pm

Hi @matthias-springer, sorry about the late reply to this thread. With recent changes in the dialect conversion drivers, is this use case supported? To recap I want to 1:N convert a subset of ops in specific op type (say load → 2 loads based on some property of the load).

I was able to do this using some hacks described above, but with recent changes in upstream this no longer works. It crashes with,

LLVM ERROR: pattern 'optimizetranspose::LoadNdOpPattern' does not support 1:N conversion

Greatly appreciate if you can help.

Thanks!

makslevental · January 21, 2025, 10:56pm

I can confirm this is doable using both a combination of close study of the framework and 1 or 2 “questionable” things; you can take a look at this PR

github.com/triton-lang/triton

[AMD] re-enable canonicalize pointers pass

triton-lang:main ← makslevental:makslevental/re-enable-canonicalize-pointers

opened 07:20AM - 04 Dec 24 UTC

makslevental

+1983 -1312

### TL;DR (too long, didn't review) This PR re-enables the `tritonamdgpu-cano…nicalize-pointers` pass[^1]. The PR is effectively a complete rewrite of the original pass, which walked the AST and mutated IR in-place, using the new [`1:N` dialect conversion framework](https://github.com/llvm/llvm-project/pull/116470). Recall a "fat pointer" is a tuple-like `(%baseptr, %offsetptr)` - the current (original) pass keeps this tuple in a global data structure while the new/rewritten pass emits this tuple into the IR as an `unrealized_cast(%baseptr, %offsetptr)`[^2]. Note, this PR also rewrites the existing lit test (see [this comment below](https://github.com/triton-lang/triton/pull/5329#issuecomment-2552629207)). ### Pass outline The pass structure/action is roughly: 1. Perform an approximate sparse dataflow analysis to find all transitive uses for `tt.func` args that are `tt.ptr`s; legalize only these ops; 2. Rewrite all operations' `use`s and `result`s to be `(%baseptr, %offsetptr)` using `ConversionPattern`s that takes the new `OneToNOpAdaptor`, which automatically forwards both `%baseptr` and `%offsetptr` through `adaptor.getOperands()`[^3]; 3. Clean up remaining `unrealized_casts` (currently only handling one category of such remaining casts but can be extended to handle all; see bullet 1 in TODOs). ### Some pre-emptive call outs Right up front I'll say this took a long time to figure out because **a)** the conversion framework is hugely complex **b)** it's being currently rewritten to be more robust/stable. As a consequence, the implementation is complex but I've tried hard to **a)** simplify as much as possible **b)** comment/note subtleties **c)** put in ample `assert`s and checks to clarify intent and gracefully fail. So some things to call out: 1. I called the dataflow analysis approximate because it does not actually use [DataFlow/SparseAnalysis](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp) and instead computes a forward slice using the heuristic "transfer function" that users of an op with a `tt.ptr` operand should be rewritten. This heuristic works because the forward slice starts from `tt.ptr` args on a `tt.func` and ends at `tt.store`, which has no results. Note, there's no reason why this component of the pass can't be a true `SparseAnalysis` implementation, it's just that this rewrite has already taken way longer than I expected (so I leave that for a possible follow-up). 2. The pass uses no global `TypeConverter` but uses local `TypeConverter`s, in `BranchInterface`/`RegionInterface` patterns. This is because **a)** we are not actually converting operand/result types (we are converting number of operands/results) **b)** the conversion framework expects/handles this lack of a `TypeConverter` exactly [the way we want](https://github.com/llvm/llvm-project/blob/399c3a78a2577c6fc68bba7f301901a0e66e87ed/mlir/lib/Transforms/Utils/DialectConversion.cpp#L1179-L1185). The local type converters are used for ops that couple to basic blocks (`^bb`s) that need to have their signatures rewritten (i.e., the ops for which we need to do `rewriter.applySignatureConversion(block, *conversion, &localTypeConverter)`). That's `scf.for`, `scf.while`, `cf.br` and `cf.cond_br` (not needed for `scf.if` which has no `bb` args). 3. `tt.func` is handled differently from all of the other ops - it is not rewritten at all. Instead, for every `%arg: tt.ptr` arg, we insert into the new body `%c0 = arith.constant 0 : i32` and `%newarg = unrealized_cast(%arg, %c0) : tt.ptr` (manually, not done by the conversion framework) and replace all uses of `%arg` by `%newarg`. These are then unpacked to `(%arg, %c0)` using `replaceOpWithMultiple` so that they "magically" appear in `adaptor.getOperands()`. Then at the end, currently, these are the only unreconciled casts (because they are the only ones **not** inserted by the conversion framework) and we materialize them by just replacing uses of `%newarg` with `%arg`. 4. `scf.if` needs to be handled specially; since it has no operands but can `yield` results, we need to rewrite it only after its `yield`s have been rewritten. This is not straightforward because the dialect conversion [does a preorder walk](https://github.com/llvm/llvm-project/blob/6ab8401f53deddbd79f930ba2bec2f824c9567e3/mlir/lib/Transforms/Utils/DialectConversion.cpp#L2705). To work around this we define legality for `scf.if` to be dependent on whether its `yield`s have been rewritten (using two `UnitAttr`s on those `yield`s). Thus, `scf.if` is "legal" and not rewritten until after the results of the `yield`s are known. ### TODOs (decisions to be made) ~~1. The new pass only reconciles the intial `unrealized_cast`s that come from `tt.func` args (see bullet 3 above). That means ops that aren't currently handled (e.g., `tt.bitcast`, `tt.scan`, etc.) will fail to be legalized and the pass will fail (e.g., [test_inductor_cummax_bool](https://github.com/triton-lang/triton/blob/main/python/test/regression/test_functional_regressions.py#L261-L263)). The old pass would just "materialize the fat pointer" for unknown operations i.e., fallback to default handling. I can certainly restore this behavior but I would prefer not to until we've done more testing that the pass is correct (i.e., that's why I've kept it behind `AMDGCN_USE_BUFFER_OPS`).~~ 2. The current pass propagates pointer attributes (such as `tt.pointer_range`) via attributes on ops. This leads to accumulating lots of attributes (such as `arith.OverflowFlags` and etc.) in the [global data structure](https://github.com/triton-lang/triton/blob/f12eae03adf741ab253a2fed64c08878f24ecc07/third_party/amd/lib/TritonAMDGPUTransforms/CanonicalizePointers.cpp#L105). I would prefer not to do this and given that we are currently discussing alternative approaches to handling such compiler hints, I did not carry over the same behavior. Again, I can certainly restore it but I propose we first decide how we would like to handle compiler hints and then afterward possibly restore it (if we decide that these attributes are how we want to implement hinting). ### Trivial TODOs Still need to add patterns for `triton::AtomicCASOp`, `triton::AtomicRMWOp`, `triton::PtrToIntOp`. ### Shoutouts Thanks @giuseros for all the advice and explanations. Shoutout @matthias-springer for both rewriting the conversion framework and advising me on how it actually works. cc @danzimm @SamGinzburg [^1]: I haven't actually moved it out of the flag but it's now usable with the `AMDGCN_USE_BUFFER_OPS` flag whereas it wasn't prior. [^2]: In reality it's the conversion framework that materializes this tuple as `unrealized_cast(%baseptr, %offsetptr)` and then reconciles/DCEs all the casts automatically. [^3]: The `unrealized_cast`s are completely "transparent" to the patterns, see [`ConversionPatternRewriterImpl::remapValues`](https://github.com/llvm/llvm-project/blob/399c3a78a2577c6fc68bba7f301901a0e66e87ed/mlir/lib/Transforms/Utils/DialectConversion.cpp#L1161).

that I recently put together (with a great deal of help from Matthias) that does exactly what you’re trying to do. It’s a fairly complex implementation so feel free to reach out if you have trouble seeing the “forest for the trees”.

charitha22 · January 21, 2025, 10:59pm

That is awesome. Thanks! I will take a look and get back to you.

I have an implementation (bit hacky) here. But I am looking for a nice way to do this.

github.com/intel/mlir-extensions

lib/Transforms/OptimizeTranspose.cpp

55753c81e


      
              // Create UnrealizedConversionCastOp to reconcile the types.
              rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(op, op.getType(),
                                                                      createNdDescOps);
          
              return success();
            }
          };
          
          // If the source tile of a LoadNdOp is split into multiple tiles due to array
          // length adjustment, We split the loadNdOp into multiple loadNdOps.
          struct LoadNdOpPattern : public OpConversionPattern<xegpu::LoadNdOp> {
            using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern;
          
            LogicalResult
            matchAndRewrite(xegpu::LoadNdOp op, OpAdaptor adaptor,
                            ConversionPatternRewriter &rewriter) const override {
              auto tdesc = adaptor.getTensorDesc();
              auto unrealizedCastOp =
                  llvm::dyn_cast_if_present<UnrealizedConversionCastOp>(
                      tdesc.getDefiningOp());
              if (!unrealizedCastOp)

matthias-springer · January 22, 2025, 7:43am

I’m not sure if it will work for your use case, but you can play with ConversionTarget: marking an op as illegal only if a conversion pattern should be applied. If an op is already legal, it will not be touched by the conversion driver.

charitha22 · January 22, 2025, 3:42pm

I am already using a conversion target. Here is my approach in summary,

I define a conversion target that filter outs the subset of ops that I need to 1:N convert using addDynamicallyLegalOp
For the ops that need 1:N conversion, I reply on UnrealizedConversionCast Op to preserve the types. Then I move on to consumers and fix them during conversion.

Code here:

github.com/intel/mlir-extensions

lib/Transforms/OptimizeTranspose.cpp

55753c81e


      
            op->walk([&](xegpu::CreateNdDescOp createNdDescOp) -> WalkResult {
              if (analysis.isCandidateForArrayLenAdjustment(createNdDescOp)) {
                markOp(createNdDescOp);
                changed = true;
              }
              return WalkResult::advance();
            });
            return changed;
          }
          
          void runArrayLengthAdjustment() {
            Operation *module = getOperation();
            TypeConverter typeConverter;
            auto addNToOneCast = [](OpBuilder &builder, Type type, ValueRange inputs,
                                    Location loc) {
              auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
              return cast.getResult(0);
            };
            typeConverter.addSourceMaterialization(addNToOneCast);
            typeConverter.addArgumentMaterialization(addNToOneCast);
            RewritePatternSet patterns(&getContext());

Tricky bit is ForOp handling. Originally I define a OneToNTypeMapping and do a signature conversion to resolve the block argument inside the ForOp. This approach is failing with recent upstream changes. For example, the block arguments are expanded after the 1:N conversion but argument materialization hook does not insert the UnrealizedConversionCast ops I need. Do you have some idea why this is happening?

Code here:

github.com/intel/mlir-extensions

lib/Transforms/OptimizeTranspose.cpp

55753c81e


      
          // Create a new ForOp with the expanded iter args.
          auto newForOp = rewriter.create<scf::ForOp>(
              op.getLoc(), op.getLowerBound(), op.getUpperBound(), op.getStep(),
              ValueRange(newIterArgs));
          // We don't need the empty blocks created by rewriter.
          rewriter.eraseBlock(newForOp.getBody());
          
          Region *region = &op.getRegion();
          Block *block = &region->front();
          // Perform signature conversion on the body block of the original forOp.
          OneToNTypeMapping signatureConverter(block->getArgumentTypes());
          
          for (size_t i = 0; i < block->getNumArguments(); i++) {
            // If iter args are expanded due to UnrealizedConversionCastOp, we need
            // 1:N type mapping in the signature converter.
            if (i >= 1 && forOpOutputTypeMapping[i - 1].second > 1) {
              auto sources =
                  adaptor.getInitArgs()[i - 1].getDefiningOp()->getOperands();
              auto sourceTypes = llvm::map_to_vector(
                  sources, [](Value source) { return source.getType(); });
              signatureConverter.addInputs(i, sourceTypes);

matthias-springer · January 22, 2025, 3:57pm

Argument materializations are deprecated. The conversion driver does not use them anymore. The callback registration function on the type converter is still there, but it will be deleted soon.

Generally speaking, any mechanism that is based on the presence of UnrealizedConversionCastOp is fragile because those are just an implementation detail of the conversion driver.

charitha22 · January 24, 2025, 6:44pm

This worked out perfectly for me. Thank you very much!

Also thanks @matthias-springer for the upstream changes. My code looks much simpler now.

Topic		Replies	Views
Request for advice: extending type conversion to 1:N conversion MLIR	10	651	June 2, 2023
Dialect conversion: type materialization MLIR	19	2605	March 10, 2021
[RFC] Merging 1:1 and 1:N Dialect Conversions MLIR	9	1170	January 4, 2025
Partial lowering with type conversions MLIR	11	1254	June 9, 2020
materializeConversion for 1:1 conversion MLIR	7	812	June 2, 2020

[Question] Conditional OneToN conversion using `OneToNOpConversionPattern`

Related topics