From e4d2bd3293da85ada6d7a19134108cac49715400 Mon Sep 17 00:00:00 2001 From: Ryan Tse Date: Sat, 30 Mar 2024 11:35:52 -0700 Subject: [PATCH] Add IR dumping to JNI. --- java/gandiva/src/main/cpp/config_builder.cc | 3 ++- java/gandiva/src/main/cpp/jni_common.cc | 24 +++++++++++++++++++ .../evaluator/ConfigurationBuilder.java | 10 ++++++-- .../arrow/gandiva/evaluator/Filter.java | 7 ++++++ .../arrow/gandiva/evaluator/JniWrapper.java | 14 +++++++++++ .../arrow/gandiva/evaluator/Projector.java | 7 ++++++ 6 files changed, 62 insertions(+), 3 deletions(-) diff --git a/java/gandiva/src/main/cpp/config_builder.cc b/java/gandiva/src/main/cpp/config_builder.cc index 85c661ee94330..3eff4adc61382 100644 --- a/java/gandiva/src/main/cpp/config_builder.cc +++ b/java/gandiva/src/main/cpp/config_builder.cc @@ -34,11 +34,12 @@ using gandiva::ConfigurationBuilder; */ JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_ConfigurationBuilder_buildConfigInstance( - JNIEnv* env, jobject configuration, jboolean optimize, jboolean target_host_cpu) { + JNIEnv* env, jobject configuration, jboolean optimize, jboolean target_host_cpu, jboolean dump_ir) { ConfigurationBuilder configuration_builder; std::shared_ptr config = configuration_builder.build(); config->set_optimize(optimize); config->target_host_cpu(target_host_cpu); + config->set_dump_ir(dump_ir); return ConfigHolder::MapInsert(config); } diff --git a/java/gandiva/src/main/cpp/jni_common.cc b/java/gandiva/src/main/cpp/jni_common.cc index a5dff9981ce89..fdac5da077044 100644 --- a/java/gandiva/src/main/cpp/jni_common.cc +++ b/java/gandiva/src/main/cpp/jni_common.cc @@ -697,6 +697,18 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_build return module_id; } +JNIEXPORT jstring JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_dumpProjectorIr( + JNIEnv* env, jobject cls, jlong module_id) { + std::shared_ptr holder = projector_modules_.Lookup(module_id); + if (holder == nullptr) { + env->ThrowNew(gandiva_exception_, "Unknown module id\n"); + return env->NewStringUTF(nullptr); + } + + std::string ir = holder->projector()->DumpIR(); + return env->NewStringUTF(ir.c_str()); +} + /// /// \brief Resizable buffer which resizes by doing a callback into java. /// @@ -977,6 +989,18 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_build return module_id; } +JNIEXPORT jstring JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_dumpFilterIr( + JNIEnv* env, jobject cls, jlong module_id) { + std::shared_ptr holder = filter_modules_.Lookup(module_id); + if (holder == nullptr) { + env->ThrowNew(gandiva_exception_, "Unknown module id\n"); + return env->NewStringUTF(nullptr); + } + + std::string ir = holder->filter()->DumpIR(); + return env->NewStringUTF(ir.c_str()); +} + JNIEXPORT jint JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_evaluateFilter( JNIEnv* env, jobject cls, jlong module_id, jint num_rows, jlongArray buf_addrs, jlongArray buf_sizes, jint jselection_vector_type, jlong out_buf_addr, diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java index fa5d285b90997..058331ed03205 100644 --- a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java +++ b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java @@ -25,10 +25,10 @@ public class ConfigurationBuilder { public long buildConfigInstance(ConfigOptions configOptions) { - return buildConfigInstance(configOptions.optimize, configOptions.targetCPU); + return buildConfigInstance(configOptions.optimize, configOptions.targetCPU, configOptions.dumpIR); } - private native long buildConfigInstance(boolean optimize, boolean detectHostCPU); + private native long buildConfigInstance(boolean optimize, boolean detectHostCPU, boolean dumpIR); public native void releaseConfigInstance(long configId); @@ -38,6 +38,7 @@ public long buildConfigInstance(ConfigOptions configOptions) { public static class ConfigOptions { private boolean optimize = true; private boolean targetCPU = true; + private boolean dumpIR = false; public static ConfigOptions getDefault() { return new ConfigOptions(); @@ -56,6 +57,11 @@ public ConfigOptions withTargetCPU(boolean targetCPU) { return this; } + public ConfigOptions withDumpIR(boolean dumpIR) { + this.dumpIR = dumpIR; + return this; + } + @Override public int hashCode() { return Objects.hash(optimize, targetCPU); diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java index 010d644d107a2..cee68c961a83a 100644 --- a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java +++ b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java @@ -185,6 +185,13 @@ private void evaluate(int numRows, List buffers, List buf } } + /** + * Dumps the IR for the LLVM module representing this filter. + */ + public String dumpIR() throws GandivaException { + return wrapper.dumpFilterIr(this.moduleId); + } + /** * Closes the LLVM module representing this filter. */ diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java index 520ef5f443e18..5f2584ed6a449 100644 --- a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java +++ b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java @@ -68,6 +68,13 @@ native void evaluateProjector(Object expander, long moduleId, int numRows, long selectionVectorBufferAddr, long selectionVectorBufferSize, long[] outAddrs, long[] outSizes) throws GandivaException; + /** + * Dumps the IR of the projector referenced by moduleId. + * + * @param moduleId moduleId to dump the IR + */ + native String dumpProjectorIr(long moduleId); + /** * Closes the projector referenced by moduleId. * @@ -111,6 +118,13 @@ native int evaluateFilter(long moduleId, int numRows, long[] bufAddrs, long[] bu int selectionVectorType, long outAddr, long outSize) throws GandivaException; + /** + * Dumps the IR of the filter referenced by moduleId. + * + * @param moduleId moduleId to dump the IR + */ + native String dumpFilterIr(long moduleId); + /** * Closes the filter referenced by moduleId. * diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java index 471ddbced0e66..4a5777f299d3f 100644 --- a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java +++ b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java @@ -350,6 +350,13 @@ private void evaluate(int numRows, List buffers, List buf outAddrs, outSizes); } + /** + * Dumps the IR for the LLVM module representing this projector. + */ + public String dumpIR() throws GandivaException { + return wrapper.dumpProjectorIr(this.moduleId); + } + /** * Closes the LLVM module representing this evaluator. */