[MLIR][DISC] legalize disc to llvm #50434

wyzero · 2021-06-24T03:08:27Z

This PR provides conversion patterns to convert disc_ral.dispatch_op
and gpu.launch_op to its disc-compatible llvm format.

This PR provides conversion patterns to convert `disc_ral.dispatch_op` and `gpu.launch_op` to its disc-compatible llvm format.

joker-eph · 2021-06-24T03:19:57Z

tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ral_legalize_to_llvm.cc

+        rewriter.getStringAttr(value), /*alignment=*/0);
+
+    rewriter.restoreInsertionPoint(ip);
+  }


Should we check that if there is an existing globalOp it also matches the provided value?

Thanks! I'll add an assert here.

joker-eph · 2021-06-24T03:21:26Z

tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ral_legalize_to_llvm.cc

+Value loadOrCreateGlobalString(PatternRewriter& rewriter, Operation* op,
+                               StringRef name, StringRef value) {
+  ModuleOp module = op->getParentOfType<ModuleOp>();
+  GlobalOp globalOp = module.lookupSymbol<GlobalOp>(name);


These lookups will be expensive, can you instead build a SymbolTable at the beginning of the function execution and update it through?
(Patterns would have to take it in their constructor and keep a reference to it)

Thanks! I'm not sure if I fully understand what you mean. I guess that you say module.lookupSymbol is expensive because this method needs to create new a SymbolTable every time it's called.

It's worth noting that RalToLLVMPass is a function-level instead module-level pass to enable maximum parallelism, thus multiple instance of RalToLLVMPass may run simultaneously. Therefore, we have to re-create a SymbolTable every time we need to lookup a symbol since other threads may insert symbols as well.

Uh that seems like a race condition: actually function-level pass can't create symbol in the outer module like that.

Now that there is no threading involved, can we look into this?

I was waiting for you here?

Sorry... I misses the updates here. I have fixed it.

joker-eph · 2021-06-24T03:39:15Z

tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ral_legalize_to_llvm.cc

+    PatternRewriter& rewriter, Operation* op) const {
+  ModuleOp module = op->getParentOfType<ModuleOp>();
+  LLVMFuncOp func = module.lookupSymbol<LLVMFuncOp>(kRalDispatchFunctionName);
+  if (!func) {


In general, prefer early exit if possible:

Suggested change

if (!func) {

if (func) return func;

joker-eph · 2021-06-24T03:46:01Z

tensorflow/compiler/mlir/hlo/tests/disc_ral_legalize_to_llvm.mlir

+    llvm.func @the_kernel() attributes {gpu.kernel} {
+      llvm.return
+    }
+  }


It seems like a precondition of this pass is that we already have a gpu.binary_blob produced. It seems to me that in such cases the module shouldn't also have an IR content?
Not only they seem redundant, but also they could get out-of-sync.

Thanks! It does not need to contain the actual IR while we still need to keep the function declaration inside the module since the gpu.launch_op needs to refer it. Thus, I'll remove the content of the the_kernel while keeping the declaration.

joker-eph · 2021-06-24T03:50:29Z

tensorflow/compiler/mlir/hlo/tests/disc_ral_legalize_to_llvm.mlir

+    // CHECK: %[[T23:.*]] = llvm.getelementptr %[[T22]]{{.*}} : (!llvm.ptr<array<90 x i8>>, i64, i64) -> !llvm.ptr<i8>
+    // CHECK: llvm.call @disc_ral_call(%[[CTX]], %[[T23]], %[[KPACK]]) : (!llvm.ptr<i8>, !llvm.ptr<i8>, !llvm.ptr<ptr<i8>>) -> ()
+
+    gpu.launch_func  @kernel_module::@the_kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) args(%0 : memref<?x?xf32>)


Can you split these patterns into different functions and use CHECK-LABEL to match the functions?

CHECK-LABEL delimit sections and makes FileCheck matching a bit more robust and debugging easier.

joker-eph · 2021-06-24T03:53:18Z

tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ral_legalize_to_llvm.cc

+class RalToLLVMPass : public RalToLLVMPassBase<RalToLLVMPass> {
+  void getDependentDialects(DialectRegistry& registry) const override {
+    registry.insert<LLVM::LLVMDialect>();
+  }


You can put this information in the .td file and it won't be needed here.

Thanks! It will introduce llvm dialect dependency into header file (e.g. PassDetail.h) if we put this dialect dependency in td file since some passes may shared one .td file.

OK, looks fine to leave it there.

wyzero · 2021-06-29T03:45:48Z

@joker-eph Hi, could you help to review the updated version? Thanks!

wyzero · 2021-06-30T01:55:56Z

@joker-eph Hi, could you help to take a look? Thanks!

wyzero · 2021-07-01T09:15:46Z

@joker-eph Hi，could you help to take a look? Thanks!

joker-eph · 2021-07-01T18:57:53Z

You may have missed my comment inline: #50434 (comment) ?

Imported from GitHub PR tensorflow/tensorflow#50434 This PR provides conversion patterns to convert `disc_ral.dispatch_op` and `gpu.launch_op` to its disc-compatible llvm format. Copybara import of the project: -- 410262f963c6ff25875355d1b7885472fc89f84e by Wenyi Zhao <reyizero@gmail.com>: [MLIR][DISC] legalize disc to llvm This PR provides conversion patterns to convert `disc_ral.dispatch_op` and `gpu.launch_op` to its disc-compatible llvm format. -- 3838e485168ce10569b3022f33225a9722073efb by Wenyi Zhao <reyizero@gmail.com>: fix -- a9e6891cfdfd317e0e7ec7001287ad87316f987d by Wenyi Zhao <reyizero@gmail.com>: fix PiperOrigin-RevId: 383282261

[MLIR][DISC] legalize disc to llvm

410262f

This PR provides conversion patterns to convert `disc_ral.dispatch_op` and `gpu.launch_op` to its disc-compatible llvm format.

google-ml-butler bot added the size:L CL Change Size: Large label Jun 24, 2021

google-ml-butler bot requested review from joker-eph, sanjoy and sherhut June 24, 2021 03:08

google-cla bot added the cla: yes label Jun 24, 2021

gbaned self-assigned this Jun 24, 2021

gbaned added this to Assigned Reviewer in PR Queue via automation Jun 24, 2021

joker-eph reviewed Jun 24, 2021

View reviewed changes

fix

3838e48

wyzero requested a review from joker-eph June 25, 2021 09:06

fix

a9e6891

joker-eph approved these changes Jul 2, 2021

View reviewed changes

google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Jul 2, 2021

PR Queue automation moved this from Assigned Reviewer to Approved by Reviewer Jul 2, 2021

kokoro-team removed the kokoro:force-run Tests on submitted change label Jul 2, 2021

copybara-service bot merged commit b033bb4 into tensorflow:master Jul 6, 2021

PR Queue automation moved this from Approved by Reviewer to Merged Jul 6, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR][DISC] legalize disc to llvm #50434

[MLIR][DISC] legalize disc to llvm #50434

wyzero commented Jun 24, 2021

joker-eph Jun 24, 2021

wyzero Jun 25, 2021 •

edited

joker-eph Jun 24, 2021

wyzero Jun 25, 2021

joker-eph Jun 29, 2021

joker-eph Jun 30, 2021

joker-eph Jul 1, 2021

wyzero Jul 2, 2021

joker-eph Jun 24, 2021

wyzero Jun 25, 2021

joker-eph Jun 24, 2021

wyzero Jun 25, 2021

joker-eph Jun 24, 2021

wyzero Jun 25, 2021

joker-eph Jun 24, 2021

wyzero Jun 25, 2021

joker-eph Jun 25, 2021

wyzero commented Jun 29, 2021

wyzero commented Jun 30, 2021

wyzero commented Jul 1, 2021

joker-eph commented Jul 1, 2021

[MLIR][DISC] legalize disc to llvm #50434

[MLIR][DISC] legalize disc to llvm #50434

Conversation

wyzero commented Jun 24, 2021

Choose a reason for hiding this comment

wyzero Jun 25, 2021 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wyzero commented Jun 29, 2021

wyzero commented Jun 30, 2021

wyzero commented Jul 1, 2021

joker-eph commented Jul 1, 2021

wyzero Jun 25, 2021 •

edited