shlo raise (#1041)

* shlo raise

* fmt

Author: wsmoses

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -152,6 +152,20 @@ def GPUWrapperOp : EnzymeXLA_Op<"gpu_wrapper", [
       OpBuilder<(ins)>];
 }
 
+def XLAWrapperOp: EnzymeXLA_Op<"xla_wrapper", [DeclareOpInterfaceMethods<SymbolUserOpInterface>, DeclareOpInterfaceMethods<CallOpInterface>, DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "XLA Call operation";
+
+  let arguments = (ins
+    SymbolRefAttr:$fn,
+    Variadic<AnyType>:$inputs
+  );
+
+  let assemblyFormat = [{
+    $fn ` ` `(` $inputs `)` attr-dict `:` functional-type($inputs, results)
+  }];
+
+}
+
 def GPUErrorOp : EnzymeXLA_Op<"gpu_error", [
   RecursiveMemoryEffects,
   SingleBlockImplicitTerminator<"enzymexla::PolygeistYieldOp">]>,

src/enzyme_ad/jax/Dialect/Ops.cpp

@@ -1426,3 +1426,48 @@ void GPUErrorOp::build(OpBuilder &builder, OperationState &result) {
   builder.createBlock(bodyRegion);
   GPUErrorOp::ensureTerminator(*bodyRegion, builder, result.location);
 }
+
+LogicalResult
+XLAWrapperOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  // TODO: Verify that the result type is same as the type of the referenced
+  // func.func op.
+  auto global = symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(
+      *this, getFnAttr());
+  if (!global)
+    return emitOpError("'")
+           << getFn() << "' does not reference a valid global funcOp";
+
+  return success();
+}
+
+void XLAWrapperOp::setCalleeFromCallable(CallInterfaceCallable callee) {
+  auto symbol = cast<SymbolRefAttr>(callee);
+  setFnAttr(cast<FlatSymbolRefAttr>(symbol));
+}
+
+CallInterfaceCallable XLAWrapperOp::getCallableForCallee() { return getFn(); }
+
+MutableOperandRange XLAWrapperOp::getArgOperandsMutable() {
+  return getInputsMutable();
+}
+
+Operation::operand_range XLAWrapperOp::getArgOperands() { return getInputs(); }
+
+ArrayAttr XLAWrapperOp::getArgAttrsAttr() { return nullptr; }
+
+void XLAWrapperOp::setArgAttrsAttr(mlir::ArrayAttr attr) { (void)attr; }
+
+ArrayAttr XLAWrapperOp::getResAttrsAttr() { return nullptr; }
+
+void XLAWrapperOp::setResAttrsAttr(ArrayAttr attr) { (void)attr; }
+
+Attribute XLAWrapperOp::removeArgAttrsAttr() { return nullptr; }
+
+Attribute XLAWrapperOp::removeResAttrsAttr() { return nullptr; }
+
+void XLAWrapperOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Read>());
+  effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Write>());
+}

src/enzyme_ad/jax/Passes/AffineToStableHLORaising.cpp

@@ -2451,6 +2451,76 @@ struct AffineToStableHLORaisingPass
       }
       funcs.pop_back();
     }
+    std::vector<enzymexla::GPUWrapperOp> gwrap;
+    op->walk([&](enzymexla::GPUWrapperOp g) { gwrap.push_back(g); });
+    for (auto g : gwrap) {
+      auto modOp = g->getParentOfType<ModuleOp>();
+      Block *body = &g->getRegion(0).front();
+      Block *newBlock = new Block();
+
+      IRMapping mapping;
+      mapping.map(body, newBlock);
+      SetVector<Value> operands;
+      getUsedValuesDefinedAbove(g->getRegion(0), operands);
+
+      SmallVector<Type> tensorTypes;
+      for (auto arg : operands) {
+        auto MT = cast<MemRefType>(arg.getType());
+        auto TT = RankedTensorType::get(MT.getShape(), MT.getElementType());
+        auto newArg = newBlock->addArgument(TT, arg.getLoc());
+        mapping.map(arg, newArg);
+        tensorTypes.push_back(TT);
+      }
+
+      auto newFuncType =
+          FunctionType::get(g->getContext(), tensorTypes, tensorTypes);
+
+      std::string name = "raised";
+
+      auto newFunc = func::FuncOp::create(g->getLoc(), name, newFuncType);
+      newFunc.setVisibility(mlir::SymbolTable::Visibility::Private);
+      newFunc.getBody().push_back(newBlock);
+
+      OpBuilder builder(newBlock, newBlock->end());
+
+      bool anyFailed = false;
+
+      llvm::DenseMap<Value, affine::AffineValueMap> maps;
+
+      ParallelContext emptyPc = ParallelContext::getEmpty(options);
+      for (auto &it : body->without_terminator()) {
+        anyFailed =
+            tryRaisingOpToStableHLO(&it, mapping, builder, maps, emptyPc)
+                .failed();
+        if (anyFailed)
+          break;
+      }
+
+      if (anyFailed) {
+        newFunc->erase();
+        return;
+      }
+
+      SmallVector<Value> results;
+      for (auto arg : operands) {
+        auto val = mapping.lookup(arg);
+        results.push_back(val);
+      }
+
+      builder.create<func::ReturnOp>(g->getLoc(), results);
+      modOp.getBody()->push_back(newFunc);
+      SymbolTable::setSymbolVisibility(newFunc,
+                                       SymbolTable::Visibility::Private);
+
+      {
+        OpBuilder builder(g);
+        auto newCall = builder.create<enzymexla::XLAWrapperOp>(
+            g->getLoc(), SymbolRefAttr::get(newFunc),
+            llvm::to_vector(operands));
+        g->erase();
+        anyRaised = true;
+      }
+    }
 
     if (!anyRaised) {
       markAllAnalysesPreserved();