feat: allow symbolref for kernel_call/jit_call (#1438)

* feat: allow symbolref for kernel_call/jit_call

* chore: run fmt

* test: use specific reactant branch

Author: avik-pal

.github/workflows/test-gb-25.yml

@@ -53,7 +53,8 @@ jobs:
           - 'main'
           # - '0123456789abcdef0123456789abcdef01234567'
         reactant_commit:
-          - 'main'
+          # - 'main'
+          - ap/symbol_ref_kernel_call
 
     steps:
       - name: Check GPUs

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -38,7 +38,7 @@ def KernelCallOp: EnzymeXLA_Op<"kernel_call", [DeclareOpInterfaceMethods<SymbolU
   let summary = "Kernel Call operation";
 
   let arguments = (ins
-    FlatSymbolRefAttr:$fn,
+    SymbolRefAttr:$fn,
     TensorI64:$gridx,
     TensorI64:$gridy,
     TensorI64:$gridz,
@@ -108,7 +108,7 @@ def JITCallOp: EnzymeXLA_Op<"jit_call", [DeclareOpInterfaceMethods<SymbolUserOpI
   let summary = "JIT Call operation";
 
   let arguments = (ins
-    FlatSymbolRefAttr:$fn,
+    SymbolRefAttr:$fn,
     Variadic<AnyType>:$inputs,
     DefaultValuedStrAttr<StrAttr, "">:$backend_config,
     OptionalAttr<AnyAttr>:$operand_layouts,

src/enzyme_ad/jax/Dialect/Ops.cpp

@@ -108,12 +108,13 @@ KernelCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
 }
 
 void KernelCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
-  auto symbol = cast<SymbolRefAttr>(callee);
-  setFnAttr(cast<FlatSymbolRefAttr>(symbol));
+  setFnAttr(cast<SymbolRefAttr>(callee));
 }
 
 CallInterfaceCallable KernelCallOp::getCallableForCallee() {
-  return SymbolRefAttr::get(getContext(), getFn());
+  auto attr = getFnAttr();
+  return SymbolRefAttr::get(getContext(), attr.getRootReference(),
+                            attr.getNestedReferences());
 }
 
 Operation::operand_range KernelCallOp::getArgOperands() { return getInputs(); }
@@ -157,8 +158,7 @@ void KernelCallOp::getEffects(
   ModuleOp moduleOp = (*this)->getParentOfType<ModuleOp>();
   assert(moduleOp && "KernelCallOp must be inside a ModuleOp");
 
-  auto callee =
-      moduleOp.lookupSymbol<FunctionOpInterface>(getFnAttr().getAttr());
+  auto callee = moduleOp.lookupSymbol<FunctionOpInterface>(getFnAttr());
   assert(callee && "KernelCallOp must have a valid function");
 
   auto effectsAttr =
@@ -184,12 +184,13 @@ LogicalResult JITCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
 }
 
 void JITCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
-  auto symbol = cast<SymbolRefAttr>(callee);
-  setFnAttr(cast<FlatSymbolRefAttr>(symbol));
+  setFnAttr(cast<SymbolRefAttr>(callee));
 }
 
 CallInterfaceCallable JITCallOp::getCallableForCallee() {
-  return SymbolRefAttr::get(getContext(), getFn());
+  auto attr = getFnAttr();
+  return SymbolRefAttr::get(getContext(), attr.getRootReference(),
+                            attr.getNestedReferences());
 }
 
 MutableOperandRange JITCallOp::getArgOperandsMutable() {
@@ -204,8 +205,7 @@ void JITCallOp::getEffects(
   ModuleOp moduleOp = (*this)->getParentOfType<ModuleOp>();
   assert(moduleOp && "JITCallOp must be inside a ModuleOp");
 
-  auto callee =
-      moduleOp.lookupSymbol<FunctionOpInterface>(getFnAttr().getAttr());
+  auto callee = moduleOp.lookupSymbol<FunctionOpInterface>(getFnAttr());
   assert(callee && "JITCallOp must have a valid function");
 
   auto effectsAttr =
@@ -1757,8 +1757,7 @@ XLAWrapperOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
 }
 
 void XLAWrapperOp::setCalleeFromCallable(CallInterfaceCallable callee) {
-  auto symbol = cast<SymbolRefAttr>(callee);
-  setFnAttr(cast<FlatSymbolRefAttr>(symbol));
+  setFnAttr(cast<SymbolRefAttr>(callee));
 }
 
 CallInterfaceCallable XLAWrapperOp::getCallableForCallee() { return getFn(); }

src/enzyme_ad/jax/Passes/LowerKernel.cpp

@@ -229,11 +229,11 @@ bool CompileGPUKernel(SymbolTableCollection &symbolTable, mlir::Location loc,
   OpBuilder rewriter(kcall);
   auto replacement = rewriter.create<enzymexla::JITCallOp>(
       kcall.getLoc(), kcall.getResultTypes(),
-      mlir::FlatSymbolRefAttr::get(kcall.getContext(), callName),
-      kcall.getInputs(), kcall.getBackendConfigAttr(),
-      kcall.getOperandLayoutsAttr(), kcall.getResultLayoutsAttr(),
-      kcall.getArgAttrsAttr(), kcall.getResAttrsAttr(),
-      kcall.getOutputOperandAliasesAttr(), kcall.getXlaSideEffectFreeAttr());
+      SymbolRefAttr::get(kcall.getContext(), callName, {}), kcall.getInputs(),
+      kcall.getBackendConfigAttr(), kcall.getOperandLayoutsAttr(),
+      kcall.getResultLayoutsAttr(), kcall.getArgAttrsAttr(),
+      kcall.getResAttrsAttr(), kcall.getOutputOperandAliasesAttr(),
+      kcall.getXlaSideEffectFreeAttr());
   kcall.replaceAllUsesWith(replacement);
   kcall.erase();
   return true;
@@ -398,11 +398,11 @@ bool CompileCPUKernel(SymbolTableCollection &symbolTable, mlir::Location loc,
   OpBuilder rewriter(kcall);
   auto replacement = rewriter.create<enzymexla::JITCallOp>(
       kcall.getLoc(), kcall.getResultTypes(),
-      mlir::FlatSymbolRefAttr::get(kcall.getContext(), callName),
-      kcall.getInputs(), kcall.getBackendConfigAttr(),
-      kcall.getOperandLayoutsAttr(), kcall.getResultLayoutsAttr(),
-      kcall.getArgAttrsAttr(), kcall.getResAttrsAttr(),
-      kcall.getOutputOperandAliasesAttr(), kcall.getXlaSideEffectFreeAttr());
+      SymbolRefAttr::get(kcall.getContext(), callName, {}), kcall.getInputs(),
+      kcall.getBackendConfigAttr(), kcall.getOperandLayoutsAttr(),
+      kcall.getResultLayoutsAttr(), kcall.getArgAttrsAttr(),
+      kcall.getResAttrsAttr(), kcall.getOutputOperandAliasesAttr(),
+      kcall.getXlaSideEffectFreeAttr());
   kcall.replaceAllUsesWith(replacement);
   kcall.erase();
   return true;

src/enzyme_ad/jax/Passes/RemoveDuplicateFuncDef.cpp

@@ -93,16 +93,6 @@ struct RemoveDuplicateFuncDefPass
         [&](SymbolUserOpInterface callOp) { callOps.push_back(callOp); });
 
     for (SymbolUserOpInterface symbolUserOp : callOps) {
-      if (auto kernelOp =
-              dyn_cast<enzymexla::KernelCallOp>(symbolUserOp.getOperation())) {
-        auto sym = kernelOp.getFn();
-        Operation *op = symbolTable.lookup(sym);
-        assert(op && "Kernel function not found");
-        auto funcOp = dyn_cast<FunctionOpInterface>(op);
-        assert(funcOp && "Kernel function is not a function");
-        if (equivalenceMap.count(funcOp.getNameAttr()))
-          kernelOp.setFn(equivalenceMap[funcOp.getNameAttr()]);
-      }
       if (auto callOp =
               dyn_cast<CallOpInterface>(symbolUserOp.getOperation())) {
         SymbolRefAttr sym = llvm::dyn_cast_if_present<SymbolRefAttr>(