[cmake][ANTLR] Support setting path to ANTLR jar (#4176)

* Support setting path to ANTLR jar

* Update comment

CMakeLists.txt

@@ -7,6 +7,7 @@ include(cmake/util/FindCUDA.cmake)
 include(cmake/util/FindVulkan.cmake)
 include(cmake/util/FindLLVM.cmake)
 include(cmake/util/FindROCM.cmake)
+include(cmake/util/FindANTLR.cmake)
 
 if(EXISTS ${CMAKE_CURRENT_BINARY_DIR}/config.cmake)
   include(${CMAKE_CURRENT_BINARY_DIR}/config.cmake)

cmake/config.cmake

@@ -144,6 +144,10 @@ set(USE_ROCBLAS OFF)
 set(USE_SORT ON)
 
 # Build ANTLR parser for Relay text format
+# Possible values:
+# - ON: enable ANTLR by searching default locations (cmake find_program for antlr4 and /usr/local for jar)
+# - OFF: disable ANTLR
+# - /path/to/antlr-*-complete.jar: path to specific ANTLR jar file
 set(USE_ANTLR OFF)
 
 # Whether use Relay debug mode

cmake/modules/ANTLR.cmake

@@ -15,29 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 if(USE_ANTLR)
-  find_program(ANTLR4 antlr4)
-
-  if (NOT ANTLR4)
-    file(GLOB_RECURSE ANTLR4JAR
-          /usr/local/lib/antlr-*-complete.jar
-          /usr/local/Cellar/*antlr-*-complete.jar)
-
-    # Get the first element of the list of antlr jars.
-    # Sort and reverse the list so the item selected is the highest
-    #   version in lib or else in Cellar if no lib installation exists.
-    list(SORT ANTLR4JAR)
-    list(REVERSE ANTLR4JAR)
-    list(GET ANTLR4JAR 0 ANTLR4JAR)
-
-    set(JAVA_HOME $ENV{JAVA_HOME})
-    if (NOT DEFINED JAVA_HOME)
-      # Hack to get system to search for Java itself.
-      set(JAVA_HOME "/usr")
-    endif()
-
-    set(ANTLR4 ${JAVA_HOME}/bin/java -jar ${ANTLR4JAR})
-  endif()
-
+  find_antlr(${USE_ANTLR})
   if(ANTLR4)
 
     set(RELAY_PARSER_DIR

cmake/util/FindANTLR.cmake

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+#######################################################
+# Enhanced version of find ANTLR.
+#
+# Usage:
+#   find_antlr(${USE_ANTLR})
+#
+# - When USE_ANTLR=ON, use auto search by first trying to find antlr4 program,
+#                      then trying to find antlr-*-complete.jar
+# - When USE_ANTLR=/path/to/antlr-*-complete.jar, use provided jar
+#
+# Provide variables:
+# - ANTLR4
+#
+macro(find_antlr use_antlr)
+  set(JAVA_HOME $ENV{JAVA_HOME})
+  if (NOT DEFINED JAVA_HOME)
+    # Hack to get system to search for Java itself.
+    message(STATUS "JAVA_HOME is not defined. Set it to ensure proper use")
+    set(JAVA_HOME "/usr")
+  endif()
+  if(MSVC)
+    set(JAVA_PROGRAM ${JAVA_HOME}/java.exe)
+  else()
+    set(JAVA_PROGRAM ${JAVA_HOME}/bin/java)
+  endif()
+  message(STATUS "Using Java at " ${JAVA_PROGRAM})
+
+  if (${use_antlr} STREQUAL "ON")
+    find_program(ANTLR4 antlr4)
+    if (NOT ANTLR4)
+      file(GLOB_RECURSE ANTLR4JAR
+          /usr/local/lib/antlr-*-complete.jar
+          /usr/local/Cellar/*antlr-*-complete.jar)
+
+      # Get the first element of the list of antlr jars.
+      # Sort and reverse the list so the item selected is the highest
+      #   version in lib or else in Cellar if no lib installation exists.
+      list(SORT ANTLR4JAR)
+      list(REVERSE ANTLR4JAR)
+      list(GET ANTLR4JAR 0 ANTLR4JAR)
+
+      set(ANTLR4 ${JAVA_PROGRAM} -jar ${ANTLR4JAR})
+    endif()
+  elseif(NOT ${use_antlr} STREQUAL "OFF")
+    set(ANTLR4 ${JAVA_PROGRAM} -jar ${use_antlr})
+  endif()
+  message(STATUS "ANTLR4="${ANTLR4})
+endmacro(find_antlr)