iree_gpu Python bindings (GPUPipelineOptionsAttr)

[makslevental]

TODO:

• map ReorderWorkgroupsStrategy (which is an enum)
• remove string match tests before landing (they were just to demo for @kuhar)

[kuhar]

Awesome, thanks for plumbing this through

[makslevental]

There's a printer ambiguity due to poor dispatch layering in AttrOrTypeFormatGen.cpp: consider

def IREEGPU_ReorderWorkgroupsStrategyAttr :
    EnumAttr<IREEGPU_Dialect, IREEGPU_ReorderWorkgroupsStrategy, "reorder_work_groups_strategy"> {
  let assemblyFormat = "$value";
  let cppNamespace = "::mlir::iree_compiler::IREE::GPU";
}

which currently generates

void ReorderWorkgroupsStrategyAttr::print(::mlir::AsmPrinter &odsPrinter) const {
  ::mlir::Builder odsBuilder(getContext());
  // shouldEmitSpace: 1
  // !lastWasPunctuation: 1
  odsPrinter << ' ';
  odsPrinter << stringifyReorderWorkgroupsStrategy(getValue());
}

(where I've inlined the values of DefFormat::shouldEmitSpace and !DefFormat::lastWasPunctuation at time of printer generation, i.e., during build).

This prints the following

// IREEGPU_ReorderWorkgroupsStrategyAttr alone
#iree_gpu<reorder_work_groups_strategy Swizzle>

// IREEGPU_ReorderWorkgroupsStrategyAttr as a Parameter to IREEGPU_GPUPipelineOptionsAttr
#iree_gpu.pipeline_options<prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false, reorder_workgroups_strategy =  Swizzle>

Note the extra space in reorder_workgroups_strategy = Swizzle.

The reason this happens is because in the latter case, the printer GPUPipelineOptionsAttr tries to do the correct thing:

void GPUPipelineOptionsAttr::print(::mlir::AsmPrinter &odsPrinter) const {
  ::mlir::Builder odsBuilder(getContext());
  odsPrinter << "<";
  {
    ...
    if (!(getReorderWorkgroupsStrategy() == ReorderWorkgroupsStrategyAttr())) {
      if (!_firstPrinted) odsPrinter << ", ";
      _firstPrinted = false;
      odsPrinter << "reorder_workgroups_strategy = ";
      if (!(getReorderWorkgroupsStrategy() == ReorderWorkgroupsStrategyAttr())) {
        // shouldEmitSpace: 0
        // !lastWasPunctuation: 0
        odsPrinter.printStrippedAttrOrType(getReorderWorkgroupsStrategy());
      }
    }
  }
  odsPrinter << ">";
}

but printStrippedAttrOrType of course has no connection to DefFormat::shouldEmitSpace and !DefFormat::lastWasPunctuation.

I believe the logical error is in pushing the responsibility of shouldEmitSpace and lastWasPunctuation on DefFormat; an attribute/token/thing should be able to print itself should be agnostic to its context I think? Here is a minimal patch that produces reasonable results in this case but I have not tested beyond:

diff --git a/mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp b/mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp
--- a/mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp	(revision a758bcdbd92efb64a3482eb95d2769d74e33f5bb)
+++ b/mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp	(date 1729175762504)
@@ -932,7 +932,7 @@
   printer.body() << "  return ::llvm::TypeSwitch<::mlir::" << valueType
                  << ", ::llvm::LogicalResult>(def)";
   const char *const printValue = R"(    .Case<{0}>([&](auto t) {{
-      printer << {0}::getMnemonic();{1}
+      printer << {0}::getMnemonic() << ' ';{1}
       return ::mlir::success();
     })
 )";
diff --git a/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp b/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
--- a/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp	(revision a758bcdbd92efb64a3482eb95d2769d74e33f5bb)
+++ b/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp	(date 1729175930071)
@@ -783,12 +783,6 @@
     os.indent();
   }
 
-  // Insert a space before the next parameter, if necessary.
-  if (shouldEmitSpace || !lastWasPunctuation)
-    os << tgfmt("$_printer << ' ';\n", &ctx);
-  shouldEmitSpace = true;
-  lastWasPunctuation = false;
-
   if (el->shouldBeQualified())
     os << tgfmt(qualifiedParameterPrinter, &ctx) << ";\n";
   else if (auto printer = param.getPrinter())

It produces the following prints (transcribing the "before"):

// before
#iree_gpu<reorder_work_groups_strategy Swizzle>
#iree_gpu.pipeline_options<prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false, reorder_workgroups_strategy =  Swizzle>
// after
#iree_gpu<reorder_work_groups_strategy Swizzle>
#iree_gpu<pipeline_options <prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false, reorder_workgroups_strategy = Swizzle>>

Note, the reason that the aggregate attribute is printed so differently is because somewhere (I have no found where...) someone is making a decision about whether it should be #iree_gpu. or #iree_gpu< based on whether a space follows the first following token (I supposed by peeking the os buffer........................).

If we are in agreement that this is a problem I can fully explore/investigate/fix.

[kuhar]

Gah. Does it make it print weird with the extra space or not work with bindings at all? @makslevental

If it's the former, I think we can live with the extra space. I've run into this before but never found the cycles to fully investigate (context)

[makslevental]

I actually meant to post/paste that whole thing to upstream.................

Gah. Does it make it print weird with the extra space or not work with bindings at all? @makslevental

It prints with the extra space. The other alternative that would tidy things up here is to change the assembly for ReorderWorkgroupsStrategyAttr to

def IREEGPU_ReorderWorkgroupsStrategyAttr :
    EnumAttr<IREEGPU_Dialect, IREEGPU_ReorderWorkgroupsStrategy, "reorder_work_groups_strategy"> {
  let assemblyFormat = "<$value>";
  let cppNamespace = "::mlir::iree_compiler::IREE::GPU";
}

which will generate prints like

#iree_gpu.pipeline_options<prefetch_shared_memory = true, no_reduce_shared_memory_bank_conflicts = false, reorder_workgroups_strategy = <Swizzle>>
#iree_gpu.reorder_work_groups_strategy<Swizzle>

I don't love it but I don't hate it either 🤷

[kuhar]

LGTM

[makslevental]

@stellaraccident if you could give us your blessing we'd appreciate it :)

[ScottTodd]

Seeing Windows build errors pointing to this: https://gist.github.com/ScottTodd/bed7de027adeafea6a7f31bf113c2284

CI has a different build error that might be triggering first, so I can't tell if this is unique to my machine/environment yet.

[ScottTodd]

Sent a PR to revert: #18833.

Also trying some things to fix: main...ScottTodd:iree:ireegpu-api-fixes

[makslevental]

Sent a PR to revert: #18833.

Also trying some things to fix: main...ScottTodd:iree:ireegpu-api-fixes

My bad - I'd seen all the generate_exports.py stuff but because this was my first such PR, it wasn't top-of-mind. Thanks @ScottTodd for reverting. If you want me to take over that branch to fix I can do that (I have a windows box I can test on).

[ScottTodd]

Sent a PR to revert: #18833.
Also trying some things to fix: main...ScottTodd:iree:ireegpu-api-fixes

My bad - I'd seen all the generate_exports.py stuff but because this was my first such PR, it wasn't top-of-mind. Thanks @ScottTodd for reverting. If you want me to take over that branch to fix I can do that (I have a windows box I can test on).

If you can iterate on it and fix it yourself, that would be great. What I have there is part of the puzzle. I still get the same error messages though. The object file linking across C APIs and Python bindings has a few tricky details and differences across platforms.