Post merge cleanup

Author: Josh Fromm

ci/jenkins/docker-images.ini

@@ -19,7 +19,7 @@
 [jenkins]
 ci_arm: tlcpack/ci-arm:20230223-070143-a3b51f11b
 ci_cortexm: tlcpackstaging/ci_cortexm:20230124-233207-fd3f8035c
-ci_cpu: tlcpack/ci-cpu:20230308-070109-9d732d0fa
+ci_cpu: tlcpack/ci-cpu:relax-20230217-001605-fcb3d9e71
 ci_gpu: tlcpack/ci-gpu:20230308-070109-9d732d0fa
 ci_hexagon: tlcpack/ci_hexagon:20230127-185848-95fa22308
 ci_i386: tlcpack/ci-i386:20221013-060115-61c9742ea

ci/jenkins/unity_jenkinsfile.groovy

@@ -32,7 +32,7 @@ import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
 // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. -->
 ci_lint = 'tlcpack/ci-lint:20221025-182121-e41d0ed6e'
 ci_gpu = 'tlcpack/ci-gpu:20221128-070141-ae4fd7df7'
-ci_cpu = 'tlcpack/ci-cpu:20230110-070003-d00168ffb'
+ci_cpu = 'tlcpackstaging/ci_cpu:relax-20230217-001605-fcb3d9e71'
 ci_wasm = 'tlcpack/ci-wasm:v0.72'
 ci_i386 = 'tlcpack/ci-i386:v0.75'
 ci_qemu = 'tlcpack/ci-qemu:v0.11'
@@ -135,9 +135,9 @@ def should_skip_ci(pr_number) {
 
 cancel_previous_build()
 
-def lint() {
+def lint(node_type) {
 stage('Prepare') {
-  node('CPU-SMALL') {
+  node(node_type) {
     // When something is provided in ci_*_param, use it, otherwise default with ci_*
     ci_lint = params.ci_lint_param ?: ci_lint
     ci_cpu = params.ci_cpu_param ?: ci_cpu
@@ -161,10 +161,12 @@ stage('Prepare') {
     """, label: 'Docker image names')
   }
 }
+}
 
+def sanity_check(node_type) {
 stage('Sanity Check') {
   timeout(time: max_time, unit: 'MINUTES') {
-    node('CPU-SMALL') {
+    node(node_type) {
       ws(per_exec_ws('tvm/sanity')) {
         init_git()
         is_docs_only_build = sh (
@@ -187,8 +189,17 @@ stage('Sanity Check') {
   }
 }
 }
+try {
+    lint('CPU-SMALL-SPOT')
+} catch(Exception ex) {
+    lint('CPU-SMALL')
+}
 
-lint()
+try {
+    sanity_check('CPU-SPOT')
+} catch(Exception ex) {
+    sanity_check('CPU')
+}
 
 // Run make. First try to do an incremental make from a previous workspace in hope to
 // accelerate the compilation. If something is wrong, clean the workspace and then
@@ -308,10 +319,8 @@ def add_hexagon_permissions() {
 // NOTE: limit tests to relax folder for now to allow us to skip some of the tests
 // that are mostly related to changes in main.
 // This helps to speedup CI time and reduce CI cost.
-stage('Build and Test') {
-  if (is_docs_only_build != 1) {
-    parallel 'BUILD: GPU': {
-      node('GPU') {
+def build_test_gpu(node_type) {
+      node(node_type) {
         ws(per_exec_ws('tvm/build-gpu')) {
           init_git()
           sh "${docker_run} ${ci_gpu} nvidia-smi"
@@ -320,16 +329,34 @@ stage('Build and Test') {
           sh "${docker_run} ${ci_gpu} ./tests/scripts/unity/task_python_relax_gpuonly.sh"
         }
       }
-    },
-    'BUILD: CPU': {
-      node('CPU-SMALL') {
+}
+
+def build_test_cpu(node_type) {
+      node(node_type) {
         ws(per_exec_ws('tvm/build-cpu')) {
           init_git()
           sh "${docker_run} ${ci_cpu} ./tests/scripts/task_config_build_cpu.sh build"
           make(ci_cpu, 'build', '-j2')
           sh "${docker_run} ${ci_cpu} ./tests/scripts/unity/task_python_relax.sh"
         }
       }
+}
+
+stage('Build and Test') {
+  if (is_docs_only_build != 1) {
+    parallel 'BUILD: GPU': {
+      try {
+        build_test_gpu('GPU-SPOT')
+      } catch(Exception ex) {
+        build_test_gpu('GPU')
+      }
+    },
+    'BUILD: CPU': {
+      try {
+        build_test_cpu('CPU-SPOT')
+    } catch(Exception ex) {
+      build_test_cpu('CPU')
+    }
     }
   } else {
     Utils.markStageSkippedForConditional('BUILD: CPU')

python/tvm/relax/frontend/__init__.py

@@ -17,4 +17,4 @@
 """
 Frontends for constructing Relax programs, with the model importers
 """
-from .common import detach_params
+from .common import detach_params, SpanContext, attach_span, emit_te_with_span

python/tvm/relax/frontend/common.py

@@ -16,9 +16,11 @@
 # under the License.
 # pylint: disable=invalid-name
 """Commons for Relax frontend."""
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union, Callable, Any
 
 import tvm
+from tvm import relax
+from ...ir import Span, SourceName
 
 
 def detach_params(mod: tvm.IRModule) -> Tuple[tvm.IRModule, Dict[str, List[tvm.nd.NDArray]]]:
@@ -53,3 +55,64 @@ def detach_params(mod: tvm.IRModule) -> Tuple[tvm.IRModule, Dict[str, List[tvm.n
         else:
             detached_mod[gv] = func
     return detached_mod, params_dict
+
+
+def emit_te_with_span(bb, func: Callable, *args: Any, **kwargs: Any) -> relax.Var:
+    """Same as block_builder.emit_te, but attaches a span to the generated call.
+    Uses the current span in the SpanContext.
+    """
+
+    call = bb.call_te(func, *args, **kwargs)
+    call = attach_span(call)
+    return bb.emit(call)
+
+
+def attach_span(op: relax.Call):
+    """Attach a span to a Relax op if it doesn't already have one.
+    Uses the current span in the SpanContext.
+    Parameters
+    ----------
+    op : relax.Expr
+        The op to attach a span to.
+    Returns
+    -------
+    op : relax.Expr
+        The op with a span attached.
+    """
+    assert isinstance(op, relax.Call), "Expected a Call node but got: {op}".format(op=str(type(op)))
+    if op.span is None:
+        return relax.Call(op.op, op.args, op.attrs, op.sinfo_args, SpanContext.current())
+    return op
+
+
+class SpanContext:
+    """A context manager for setting the current Span.
+    Parameters
+    ----------
+    span : Union[Span, str]
+        The span to set as the current span.
+    """
+
+    __current_span = None
+
+    def __init__(self, span: Union[Span, str]):
+        assert isinstance(span, (Span, str)), "span must be a Span or str"
+        if isinstance(span, str):
+            span = Span(SourceName(span), 0, 0, 0, 0)
+        SpanContext.__current_span = span
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, ptype, value, trace):
+        SpanContext.__current_span = None
+
+    @staticmethod
+    def current():
+        """Get the span in the current context.
+        Returns
+        -------
+        span : Optional[Span]
+            The current span.
+        """
+        return SpanContext.__current_span

python/tvm/relax/transform/legalize_ops/statistical.py

@@ -51,9 +51,17 @@ def _te_variance(x: te.Tensor, axis: List[tir.IntImm], keepdims: bool) -> te.Ten
 
 @register_legalize("relax.mean")
 def _mean(bb: BlockBuilder, call: Call) -> Expr:
-    return bb.call_te(
-        _te_mean, call.args[0], call.attrs.axis, call.attrs.keepdims, primfunc_name_hint="mean"
+    # Reductions often overflow fp16 values. If we encounter one, just cast to fp32.
+    data = call.args[0]
+    original_dtype = call.args[0].struct_info.dtype
+    if original_dtype == "float16":
+        data = bb.normalize(bb.emit_te(topi.cast, data, "float32"))
+    output = bb.normalize(
+        bb.call_te(_te_mean, data, call.attrs.axis, call.attrs.keepdims, primfunc_name_hint="mean")
     )
+    if output.struct_info.dtype != original_dtype:
+        output = bb.emit_te(topi.cast, output, original_dtype)
+    return output
 
 
 @register_legalize("relax.std")
@@ -68,13 +76,23 @@ def te_std(x: te.Tensor, axis: List[tir.IntImm], keepdims: bool) -> te.Tensor:
 
 @register_legalize("relax.variance")
 def _variance(bb: BlockBuilder, call: Call) -> Expr:
-    return bb.call_te(
-        _te_variance,
-        call.args[0],
-        call.attrs.axis,
-        call.attrs.keepdims,
-        primfunc_name_hint="variance",
+    # Reductions often overflow fp16 values. If we encounter one, just cast to fp32.
+    data = call.args[0]
+    original_dtype = call.args[0].struct_info.dtype
+    if original_dtype == "float16":
+        data = bb.normalize(bb.emit_te(topi.cast, data, "float32"))
+    output = bb.normalize(
+        bb.call_te(
+            _te_variance,
+            data,
+            call.attrs.axis,
+            call.attrs.keepdims,
+            primfunc_name_hint="variance",
+        )
     )
+    if output.struct_info.dtype != original_dtype:
+        output = bb.emit_te(topi.cast, output, original_dtype)
+    return output
 
 
 register_legalize("relax.max", _statistical(topi.max))

python/tvm/relax/vm_build.py

@@ -293,7 +293,8 @@ def foo(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")):
         target = tvm.target.Target(target)
 
     passes = []
-    passes.append(relax.transform.RewriteDataflowReshape())
+    # TODO(jwfromm) Reenable once slice bug is fixed.
+    # passes.append(relax.transform.RewriteDataflowReshape())
     passes.append(relax.transform.ToNonDataflow())
     passes.append(relax.transform.CallTIRRewrite())
     passes.append(relax.transform.StaticPlanBlockMemory())

src/relax/ir/block_builder.cc

@@ -574,7 +574,7 @@ class Normalizer : public BlockBuilderImpl, private ExprFunctor<Expr(const Expr&
     if (unchanged) {
       call = GetRef<Call>(op);
     } else {
-      call = Call(new_op, new_args, op->attrs, op->sinfo_args);
+      call = Call(new_op, new_args, op->attrs, op->sinfo_args, op->span);
     }
 
     if (!call->struct_info_.defined()) {

tests/scripts/unity/task_python_relax.sh

@@ -29,7 +29,9 @@ export TVM_NUM_THREADS=2
 make cython3
 
 # Run Relax tests
-TVM_TEST_TARGETS="${TVM_RELAY_TEST_TARGETS:-llvm}" pytest tests/python/relax
+export TVM_TEST_TARGETS="${TVM_RELAY_TEST_TARGETS:-llvm}"
+export PLATFORM=cpu
+run_pytest cython unity-relax tests/python/relax
 
 # Run Relax examples
 # python3 ./apps/relax_examples/mlp.py