[TIR][VM] Revert a change to lower_tvm_builtin.cc from #6126

[mbrookhart]

This change is causing a regression in tests/python/fronend/onnx/test_forward.py:test_loop, and causes the generated IR for the shape function to be very different.

I'm really, really confused why this didn't fail in CI, it fails in every local setup and CI docker image we've tried outside of CI.

cc @jwfromm @tmoreau89 @tqchen @zhanghaohit

this PR:

allocate(v_copy_shape_func: Pointer(int64), int64, [1]) {
  attr [0] "extern_scope" = 0 {
    v_expand_dim_shape_func: Pointer(int64)[0] = 1i64
    v_expand_dim_shape_func[1] = (int64*)v_copy_shape_func[0]
  }
  attr [0] "extern_scope" = 0 {
    v_concatenate_shape_func: Pointer(int64)[0] = 0i64
    v_concatenate_shape_func[0] = ((int64*)v_concatenate_shape_func[0] + (int64*)placeholder: Pointer(int64)[0])
    v_concatenate_shape_func[0] = ((int64*)v_concatenate_shape_func[0] + (int64*)v_expand_dim_shape_func[0])
    v_concatenate_shape_func[1] = (int64*)placeholder[1]
    assert(((int64*)v_concatenate_shape_func[1] == (int64*)v_expand_dim_shape_func[1]), "Dims mismatch in the inputs of concatenate.")
    0
  }
}

current main:

attr [v_expand_dim_shape_func: Pointer(int64)] "storage_alignment" = 128 {
  let v_expand_dim_shape_func = @tir.TVMBackendAllocWorkspace(1, dev_id: int32, 16u64, 0, 64, dtype=handle)
   {
    if @tir.isnullptr(v_expand_dim_shape_func, dtype=bool) {
      @tir.tvm_throw_last_error(, dtype=int32)
    }
    attr [v_copy_shape_func: Pointer(int64)] "storage_scope" = "global";
    attr [v_copy_shape_func] "storage_alignment" = 128 {
      let v_copy_shape_func = @tir.TVMBackendAllocWorkspace(1, dev_id, 8u64, 0, 64, dtype=handle)
       {
        if @tir.isnullptr(v_copy_shape_func, dtype=bool) {
          @tir.tvm_throw_last_error(, dtype=int32)
        }
         {
          attr [0] "extern_scope" = 0 {
            v_expand_dim_shape_func[0] = 1i64
            v_expand_dim_shape_func[1] = (int64*)v_copy_shape_func[0]
          }
          attr [0] "extern_scope" = 0 {
            v_concatenate_shape_func: Pointer(int64)[0] = 0i64
            v_concatenate_shape_func[0] = ((int64*)v_concatenate_shape_func[0] + (int64*)placeholder: Pointer(int64)[0])
            v_concatenate_shape_func[0] = ((int64*)v_concatenate_shape_func[0] + (int64*)v_expand_dim_shape_func[0])
            v_concatenate_shape_func[1] = (int64*)placeholder[1]
            assert(((int64*)v_concatenate_shape_func[1] == (int64*)v_expand_dim_shape_func[1]), "Dims mismatch in the inputs of concatenate.")
            0
          }
        }
      }
      if (@tir.TVMBackendFreeWorkspace(1, dev_id, v_copy_shape_func, dtype=int32) != 0) {
        @tir.tvm_throw_last_error(, dtype=int32)
      }
    }
  }
  if (@tir.TVMBackendFreeWorkspace(1, dev_id, v_expand_dim_shape_func, dtype=int32) != 0) {
    @tir.tvm_throw_last_error(, dtype=int32)
  }
}

[jwfromm]

I guess unsurprisingly this revision causes the error that lead to the code being removed in #6126. We'll have to either fix it a different way or temporarily disable that portion of VTA testing.

[tmoreau89]

We'll want to not disable the VTA tests because that backend will quickly bitrot.

At this point, we can either:
(1) Revert #6126 but that is less than desirable
(2) Identify a fix that resolves the tests/python/fronend/onnx/test_forward.py:test_loop and VTA unit test - but this can take a couple days

[tmoreau89]

I tried to reproduce the failing VTA test with the changes applied by @mbrookhart but I wasn't able to reproduce the error https://ci.tlcpack.ai/blue/organizations/jenkins/tvm/detail/PR-8274/1/pipeline#step-458-log-1910. To be continued.

[zhanghaohit]

In my option, even though this revert makes tests/python/fronend/onnx/test_forward.py:test_loop pass, there may be a hidden bug for the vm runtime/other implementation. What if the condition if (constant_size > 0 && constant_size * nbytes < runtime::kMaxStackAlloca) is not true, it will generate the same IR as the current main.

I admit that the removal of the following optimization:

    if (device_type_.defined()) {
      if (const auto* dev_type = device_type_.as<IntImmNode>()) {
        if (dev_type->value == kDLCPU) {
          int32_t constant_size = op->constant_allocation_size();
          if (constant_size > 0 && constant_size * nbytes < runtime::kMaxStackAlloca) {
            return stmt;
          }
        }
      }
    }

may cause some performance regressions. But this optimization will raise LLVM function signature errors when there are multiple targets (For now, I cannot find an alternative fix to make this work).

So there may be two bugs in the current codebase:

• the VM runtime (also maybe other parts) cannot handle the IR generated by the current main (e.g., if constant_size > runtime::kMaxStackAlloca)
• when there are multiple targets, the generated IR / LLVM code with this PR is not correct.

If this problem blocks something urgent to merge, I think we can do a quick fix first (choose either one of the following two) and open a bug issue for further fix:

• to fix the VTA unittest, we can remove the multiple targets test in deploy_classification.py. Just remove the "sim" in these three lines: deploy_classification.py#L194, deploy_classification.py#L206
  deploy_classification.py#L224
• to quick fix the test tests/python/fronend/onnx/test_forward.py:test_loop, remove the check assert out[i] == inputs[j][i], "Dims mismatch in the inputs of concatenate." in the _concatenate_shape_func.

For the two hidden bugs, I think it takes some time to find and fix them.

What do you think?

[mbrookhart]

ping @tmoreau89 @zhanghaohit any update on this?

[mbrookhart]

@tmoreau89 @zhanghaohit

I took this suggestion:

to fix the VTA unittest, we can remove the multiple targets test in deploy_classification.py. Just remove the "sim" in these three lines: deploy_classification.py#L194, deploy_classification.py#L206
deploy_classification.py#L224

To disable the tutorial.

[mbrookhart]

@tmoreau89 @zhanghaohit those comments did not, in fact, prevent the tutorial from being run, so this still failed. I'm fixing the CI issues onnx tests today, and this is now blocking that. Any thoughts on how to proceed?

[zhanghaohit]

See this: #8643

@tmoreau89 @zhanghaohit those comments did not, in fact, prevent the tutorial from being run, so this still failed. I'm fixing the CI issues onnx tests today, and this is now blocking that. Any thoughts on how to proceed?

[zhanghaohit]

@tmoreau89 @zhanghaohit those comments did not, in fact, prevent the tutorial from being run, so this still failed. I'm fixing the CI issues onnx tests today, and this is now blocking that. Any thoughts on how to proceed?

@mbrookhart I think maybe you apply the changes from this PR (#8643) to see whether it works, since we already discuss a lot here. I will delete #8643 later.

[tmoreau89]

Thanks @mbrookhart and @zhanghaohit the fix has been merged.

[tmoreau89]

@zhanghaohit would you be able to create an issue that tracks the two hidden bugs that require deeper investigation?

[zhanghaohit]

@zhanghaohit would you be able to create an issue that tracks the two hidden bugs that require deeper investigation?

Sure. I create two issues: #8977 #8978

[manupak]

Hi, @mbrookhart

Can you explain the rationale for not generating TVMBAW calls for static size allocates for DLCPU ?

cc: @grant-arm @Mousius @mbaret

[manupak]

@tmoreau89 @areusch @mbrookhart ,

IIUC,

This is problematic for micro because now the constant sized allocates are forced to be placed on stack (bypassing TVMPlatformAllocate abstraction) because that is how the codegen_c lowers the allocate.

tvm/src/target/source/codegen_c.cc

Lines 860 to 877 in 1b99adc

	void CodeGenC::VisitStmt_(const AllocateNode* op) {
	ICHECK(!is_zero(op->condition));
	std::string vid = AllocVarID(op->buffer_var.get());
	this->PrintIndent();
	int32_t constant_size = op->constant_allocation_size();
	ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now";
	auto scope = GetPtrStorageScope(op->buffer_var);
	alloc_storage_scope_[op->buffer_var.get()] = scope;
	PrintStorageScope(scope, stream);
	PrintType(op->dtype, stream);
	stream << ' ' << vid << '[' << constant_size << "];\n";
	RegisterHandleType(op->buffer_var.get(), op->dtype);
	this->PrintStmt(op->body);
	}

I dont think is desired.

[mbrookhart]

Hi @manupa-arm I think we have some conflict between what the VM needs to run on CPU and what microprocessors need. Can we talk about this on the issues @zhanghaohit created #8977 #8978?

I think @jroesch might have some thoughts on separating the needs.