This is an archive of the discontinued LLVM Phabricator instance.

[MLIR][mlir-spirv-cpu-runner] A pass to emulate a call to kernel in LLVM
ClosedPublic

Authored by georgemitenkov on Aug 17 2020, 2:51 PM.

Details

Summary

This patch introduces a pass for running
mlir-spirv-cpu-runner - LowerHostCodeToLLVMPass.

This pass emulates gpu.launch_func call in LLVM dialect and lowers
the host module code to LLVM. It removes the gpu.module, creates a
sequence of global variables that are later linked to the varables
in the kernel module, as well as a series of copies to/from
them to emulate the memory transfer to/from the host or to/from the
device sides. It also converts the remaining Standard dialect into
LLVM dialect, emitting C wrappers.

Diff Detail

Event Timeline

georgemitenkov created this revision.Aug 17 2020, 2:51 PM
georgemitenkov requested review of this revision.Aug 17 2020, 2:51 PM
mravishankar requested changes to this revision.Aug 17 2020, 11:37 PM

I did an initial pass on this. Will come back and look in more detail, but have some high-level comments.

mravishankar added inline comments.Aug 17 2020, 11:37 PM
mlir/include/mlir/Conversion/Passes.td
94

Looking through this, I am not sure why this is split into two. It might be better to combine this into one pass. AFAICS you have all the information you need to convert a gpu.launch_func to

  • a series of copies from argument to the global variables and then the
  • Actual call to the kernel function
  • a series of a copies to copy the result back.
mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
181

The GPU Dialect adds the gpu.container_module to identify kernel side. Further the LLVM -> SPIR-V lowering preserve the spv.target_env attribute. Its better to check for one or both of these attributes to be more targeted in what is treated as a kernel module?

This revision now requires changes to proceed.Aug 17 2020, 11:37 PM
georgemitenkov added inline comments.Aug 18 2020, 1:04 AM
mlir/include/mlir/Conversion/Passes.td
94

The reason to split comes from the fact that we cannot convert Standard to LLVM before our pass (Then gpu.launch_func arguments' types do not match as memrefs are converted into memref descriptors in LLVM). So we would have something like this:

"gpu.launch_func"(%24, %24, %24, %24, %24, %24, %16) {kernel = @kernels::@simple} : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.i64, !llvm.struct<(ptr<i32>, ptr<i32>, i64, array<1 x i64>, array<1 x i64>)>) -> ()

Nor we can run std-to-llvm after. If we replace the launch op with llvm.call, then its arguments will also be of LLVMType. This means we cannot really write the series of copies from argument to the global variables. This is because arguments' types are in LLVM but actual values passed are still in standard.

My solution to this was to preprocess gpu.launch_func so that it can be lowered to LLVM as a normal function, and then replace it as intended when everything is in LLVM. There might be something I am missing here, I will have a look at how to combine 2 passes in one (which is a definitely a nicer way) :)

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
181

Thanks! I am not sure how gpu.container_module in general helps because it only identifies that it contains a kernel module, but may also have other modules in it?
In this case, we can check for the module that is != gpu.container_module as we consider only 2 modules.

I think reusing spv.target_env and spv.module's attributes is a great way, but for that we have to have a way of passing this info in SPIR-V to LLVM conversion. Something like

module attires { "kernel" } { ...

Also, I am wondering why GPU-> SPIR-V module conversion doesn't pass the symbolic name? If we would have a SPIR-V module with the name "spv.{gpu module name here}" for example, we can find the kernel straight away by looking up in symbol table?

georgemitenkov marked 2 inline comments as not done.Aug 18 2020, 1:06 AM
georgemitenkov marked an inline comment as done.Aug 19 2020, 8:59 AM
georgemitenkov added inline comments.
mlir/include/mlir/Conversion/Passes.td
94

Actually, there is a way to do it by pulling in std-to-llvm patterns. I will update the patch with a new combined version of the passes.

Combined 2 passes into 1.
Finding kernel module is the same for now (to be changed).

Added a utilty function for kernel global variable lookup.

georgemitenkov marked an inline comment as done.Aug 20 2020, 5:19 AM

Added missing change.

Harbormaster completed remote builds in B69016: Diff 286787.
georgemitenkov edited the summary of this revision. (Show Details)Aug 20 2020, 7:00 AM
georgemitenkov retitled this revision from [MLIR][mlir-spirv-runner] Passes for spirv-runner to [MLIR][mlir-spirv-cpu-runner] Passes for spirv-cpu-runner.Aug 20 2020, 7:10 AM
mravishankar requested changes to this revision.Aug 20 2020, 10:59 PM
mravishankar added inline comments.
mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
111

Correct me if I am wrong, but it seems like the assumption here is that the structure is actually

module {
   spv.module { 
      spv.func @foo 
   }
   module @kernel_module_name {gpu.container_module} {
      gpu.func @foo
   }
   
   gpu.launch_func @kernel_module_name::foo
}

This is the case now, but that could change. We should be really looking at at the spv.module directly. There some steps though to get there

  1. You probably need to add SymAttr to spv.module. As you mentioned, that when converting the gpu module to spir-v module, we can use the same symbol name
  2. spv.EntryPoint actually is supposed to have a list of spv.globalVariables that are accessed within the entry point. We dont emit it currently cause it isnt needed by Vulkan side. But that is useful here to get the "arguments". The spv.EntryPoint is added by the LowerABIAttributes pass. Probably need to do it when lowering a func to spv.func where you have all the information.

This is a fairly big change though. A intermediate stage would be given that you check for a single spv.Module and a single entry point function, under that assumption, all non builtin spv.globalVariables are kernel arguments, and they have the [binding,descriptor_set] as [0, 0], [1,0], [2, 0]... and so on. So argument 0 -> [0, 0], argument 1 -> [1, 0], etc. Use that to define/find the symbol that is used to do the copy. Does this make sense.

146

Can we also append the gpu.container_module name and the kernel function name to avoid conflicts.

181

Also, I am wondering why GPU-> SPIR-V module conversion doesn't pass the symbolic name? If we would have a SPIR-V module with the name "spv.{gpu module name here}" for example, we can find the kernel straight away by looking up in symbol table?

Good point. We didnt need it, but totally agree that the spv.module should have a symbol name as well. (See comment above).

This revision now requires changes to proceed.Aug 20 2020, 10:59 PM
georgemitenkov added inline comments.Aug 21 2020, 2:10 AM
mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
111

Actually the current structure was

module {gpu.container_module} {
  module {
    llvm.func() @foo
  }
  gpu.module @kernel_module_name {
      gpu.func @foo
  }

  gpu.launch_func @kernel_module_name::foo
}

because we first convert SPIR-V to LLVM for easier global variable type deduction. But dealing with gpu.launch_func op before makes more sense as we do not lose spv.EntryPoint and binding/descriptor set info.

Probably need to do it when lowering a func to spv.func where you have all the information.

I think that this can be done within LowerABIAttributes where we convert spv.func(args) to spv.func()?

Then the ideal flow would be:

  • look up the needed SPIR-V module by __spv_{kernel_module_name}
  • look up an entry point for this kernel and get the list of spv.globalVariables associated with this kernel. I think these globals are named by
auto name = funcOp.getName().str() + "_arg_" + std::to_string(argIndex);

I think it will be better to append the module name in the beginning as well?

  • Create the llvm.global (need to convert the type here) with the name:
name + "_binding0_descriptor_set_1"

The same naming convention can be used when we lower SPIR-V global's to LLVM. This would remove the need of EncodeDescriptorSetsPass then cause it can be done in ConvertSPIRVToLLVM?

georgemitenkov marked 3 inline comments as done.

This is a big update to the pass structure. Now it is run after GPU to SPIR-V
conversion (and all ABI lowering, etc.). This allows to use SPIR-V module,
binds of global variables for kernel call emulation. Also, we can have multiple
kernels in the program (To fully run these the linking has to be updated).

Also, kernel function and kernel operands (global variables) are renamed during
the pass to have a kernel name in front of their names.

EncodeDescriptorSetsPass is removed as binding encoding can be done in SPIR-V
to LLVM conversion (TODO).

Added missing file.

Removed EncodeDescriptorSetsPass.

mravishankar added inline comments.Aug 24 2020, 12:09 PM
mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
111

Thanks for clarifying. This makes it clear whats happening.

The LowerABIAttributes pass already does what you want for the spv.globalVariable (here). I think encoding the name as you mentioned seems fine to me. Removing an extra pass is good as well.

Rebase on top of other patches and minor changes: variables/function names now keep __spv__ from SPIR-V module.

georgemitenkov marked an inline comment as done.Aug 26 2020, 10:39 AM
georgemitenkov retitled this revision from [MLIR][mlir-spirv-cpu-runner] Passes for spirv-cpu-runner to [MLIR][mlir-spirv-cpu-runner] A pass to emulate a call to kernel in LLVM.Aug 26 2020, 10:41 AM
georgemitenkov edited the summary of this revision. (Show Details)
rriddle added inline comments.Aug 26 2020, 2:57 PM
mlir/include/mlir/Conversion/Passes.td
95

Is this wrapped at 80 characters?

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
98

Please avoid using distance as it is O(N) in size. You can use llvm::hasSingleElement() for this check.

99

nit: You can return this directly, it auto-converts to failure().

115

nit: Please just use module.sym_name() instead.

127

You can set an attribute via module.sym_name(newName).

254

nit: Spell out auto.

290

nit: Drop trivial braces.

georgemitenkov marked 6 inline comments as done.Aug 26 2020, 8:29 PM
georgemitenkov added inline comments.
mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
127

I am not sure this can be done like that? Do you mean using sym_nameAttr() and passing a string attribute instead?

Addressed comments.

mravishankar requested changes to this revision.Aug 26 2020, 9:46 PM

This looks good! I just have a few final nits.

Btw, It would be better for me to commit this for you. Please let me know when this is ready, and I will commit this. I am setting this as "request change" only so that once you update this it shows up on my dashboard.

Thanks for addressing all the issues raised! It looks fairly clean to me now.

mlir/lib/Conversion/GPUCommon/ConvertLaunchFuncToLLVMCalls.cpp
30

Nit: Instead of macros I would rather have a static method that returns a std::string.

162

Nit: Just add { ... } around statements that span multiple lines...

177

Use OpBuilder::InsertionGuard(rewriter). It will reset the insertion point after the scope.

This revision now requires changes to proceed.Aug 26 2020, 9:46 PM
georgemitenkov marked 3 inline comments as done.

Addressed comments and rebased on master to pick up committed patches.

Fixed StringRef error.

Added a missing change.

flaub added a subscriber: flaub.Aug 30 2020, 1:45 PM
mravishankar accepted this revision.Oct 14 2020, 11:01 PM

Coming around to landing this change now. This was causing some issues with builds (internal builds within Google) and I figured a better placement of the files works better. See this patch : https://reviews.llvm.org/D89448 for the changes. I will commit that change instead. The dependent change should be modified to suit this though. I can take care of that as well

This revision is now accepted and ready to land.Oct 14 2020, 11:01 PM

Moved files around as suggested by @mravishankar.

Coming around to landing this change now. This was causing some issues with builds (internal builds within Google) and I figured a better placement of the files works better. See this patch : https://reviews.llvm.org/D89448 for the changes. I will commit that change instead. The dependent change should be modified to suit this though. I can take care of that as well

Great, thanks! I will update the dependent runner patch accordingly.

Updated struct types according to https://reviews.llvm.org/D87206.

@mravishankar So this patch should be good to land then?

This revision was landed with ongoing or failed builds.Oct 26 2020, 5:14 AM
This revision was automatically updated to reflect the committed changes.