use discardable attributes

bytedance · Sep 6, 2024 · deac541 · deac541
1 parent e114514
commit deac541
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 9 deletions.
diff --git a/compiler/lib/Conversion/FuncToByre/FuncToByre.cpp b/compiler/lib/Conversion/FuncToByre/FuncToByre.cpp
@@ -108,13 +108,6 @@ class ConvertGPULaunchFuncToByrePattern
         launchOp->getLoc(), TypeRange(), "PTXOp", launchOp.getKernelOperands(),
         /*memEffects*/ ArrayAttr());
 
-    computeOp->setAttr(rewriter.getStringAttr("device_file_name"),
-                       launchOp->getAttr("device_file_name"));
-    if (launchOp->hasAttr(byre::getKernelCallConventionAttrName())) {
-      computeOp->setAttr(
-          byre::getKernelCallConventionAttrName(),
-          launchOp->getAttr(byre::getKernelCallConventionAttrName()));
-    }
     computeOp->setAttr(
         rewriter.getStringAttr("kernel_name"),
         rewriter.getStringAttr(launchOp.getKernelName().getValue()));
@@ -135,6 +128,11 @@ class ConvertGPULaunchFuncToByrePattern
     computeOp->setAttr("BlockSize.y", rewriter.getI32IntegerAttr(by));
     computeOp->setAttr("BlockSize.z", rewriter.getI32IntegerAttr(bz));
 
+    // get discardable attributes from gpu.launch
+    for (auto attr : launchOp->getDiscardableAttrDictionary().getValue()) {
+      computeOp->setAttr(attr.getName(), attr.getValue());
+    }
+
     rewriter.eraseOp(launchOp);
 
     return success();

diff --git a/compiler/python/byteir/compile.py b/compiler/python/byteir/compile.py
@@ -131,7 +131,6 @@ def _compile_cuda(
             PassManager.parse("builtin.module(gpu-opt{device-file-name=" + output_file_prefix + ".ptx" + "})").run(module.operation)
         _print_verbose(module, "// IR Dump After GPU Opt:") if verbose else ...
     with context:
-        # PassManager.parse("builtin.module(func.func(remove-func-body{anchor-attr=__byteir_elementwise_fusion__}))").run(module.operation)
         PassManager.parse("builtin.module(inline)").run(module.operation)
         PassManager.parse("builtin.module(func.func(lccl-to-byre))").run(module.operation)
         PassManager.parse("builtin.module(func.func(gpu-launch-func-to-byre))").run(module.operation)
@@ -238,7 +237,6 @@ def _compile_cuda_with_ait(
             PassManager.parse("builtin.module(gpu-opt{device-file-name=" + output_file_prefix + ".ptx" + "})").run(module.operation)
         _print_verbose(processor.module, "// IR Dump After GPU Opt:") if verbose else ...
     with context:
-        # PassManager.parse("builtin.module(func.func(remove-func-body{anchor-attr=__byteir_elementwise_fusion__}))").run(processor.module.operation)
         PassManager.parse("builtin.module(inline)").run(processor.module.operation)
         PassManager.parse("builtin.module(func.func(lccl-to-byre))").run(module.operation)
         PassManager.parse("builtin.module(func.func(gpu-launch-func-to-byre))").run(processor.module.operation)